diff --git a/cli/openbb_cli/argparse_translator/argparse_class_processor.py b/cli/openbb_cli/argparse_translator/argparse_class_processor.py index e2e4f3cd48f..3095c8fec7d 100644 --- a/cli/openbb_cli/argparse_translator/argparse_class_processor.py +++ b/cli/openbb_cli/argparse_translator/argparse_class_processor.py @@ -68,11 +68,15 @@ class ArgparseClassProcessor: @classmethod def _custom_groups_from_reference(cls, class_name: str, function_name: str) -> dict: - route = f"/{class_name.replace('_', '/')}/{function_name}" - reference = {route: cls._reference[route]} if route in cls._reference else {} - if not reference: + # Match by flattened route, since a namespace name may contain ``_``. + target = f"{class_name}_{function_name}" + route = next( + (r for r in cls._reference if r.strip("/").replace("/", "_") == target), + None, + ) + if route is None: return {} - rp = ReferenceToArgumentsProcessor(reference) + rp = ReferenceToArgumentsProcessor({route: cls._reference[route]}) return rp.custom_groups.get(route, {}) # ty: ignore[invalid-return-type] @classmethod diff --git a/cli/openbb_cli/cli.py b/cli/openbb_cli/cli.py index e4ec809441e..ee675e8b7db 100644 --- a/cli/openbb_cli/cli.py +++ b/cli/openbb_cli/cli.py @@ -591,7 +591,11 @@ def _run_spec_one_shot( ) -> int: """Dispatch a single command using one or more precomputed ``.spec`` files.""" from openbb_cli.dispatchers.protocol import Request, Response - from openbb_cli.dispatchers.runtime import _to_json_line + from openbb_cli.dispatchers.runtime import ( + _follow_stream, + _looks_like_stream, + _to_json_line, + ) from openbb_cli.dispatchers.spec import SpecCommandError, parse_command_argv if not command_argv: @@ -622,8 +626,13 @@ def _run_spec_one_shot( await dispatcher.aclose() response: Response = asyncio.run(_dispatch_and_close()) + + if response.ok and _looks_like_stream(response.result): + return _follow_stream(response.result) + sys.stdout.write(_to_json_line(response) + "\n") sys.stdout.flush() + return 0 if response.ok else 1 @@ -859,7 +868,11 @@ def _run_spec_dispatch( ) -> int: """Schema-validated one-shot dispatch reusing an existing dispatcher.""" from openbb_cli.dispatchers.protocol import Request, Response - from openbb_cli.dispatchers.runtime import _to_json_line + from openbb_cli.dispatchers.runtime import ( + _follow_stream, + _looks_like_stream, + _to_json_line, + ) from openbb_cli.dispatchers.spec import SpecCommandError, parse_command_argv try: @@ -872,8 +885,13 @@ def _run_spec_dispatch( return await dispatcher.dispatch(Request(command=command, params=params)) response: Response = asyncio.run(_dispatch()) + + if response.ok and _looks_like_stream(response.result): + return _follow_stream(response.result) + sys.stdout.write(_to_json_line(response) + "\n") sys.stdout.flush() + return 0 if response.ok else 1 diff --git a/cli/openbb_cli/codegen/fetcher_gen.py b/cli/openbb_cli/codegen/fetcher_gen.py index fb2b4626b22..133df592d7e 100644 --- a/cli/openbb_cli/codegen/fetcher_gen.py +++ b/cli/openbb_cli/codegen/fetcher_gen.py @@ -76,6 +76,30 @@ class GeneratedFetcher: fetcher_class: str source: str credentials_used: dict[str, dict[str, str]] = field(default_factory=dict) + is_streaming: bool = False + + +def _command_is_streaming(cmd_spec: dict[str, Any]) -> bool: + """Return whether a command's success response is a stream. + + Parameters + ---------- + cmd_spec : dict + The spec command entry. + + Returns + ------- + bool + True when a 2xx response advertises ``text/event-stream``. + """ + for status, content in (cmd_spec.get("response_schemas") or {}).items(): + if ( + str(status).startswith("2") + and isinstance(content, dict) + and any("event-stream" in str(ct) for ct in content) + ): + return True + return False def _module_name_from_command(name: str) -> str: @@ -543,6 +567,32 @@ def generate_fetcher_module(spec: FetcherCommandSpec) -> GeneratedFetcher: ) header_block = _header_dict_construction(creds) + if _command_is_streaming(spec.cmd_spec): + source = _render_streaming_fetcher_source( + spec=spec, + qp_class=qp_class, + qp_classes=qp_generated.flatten(), + data_class=data_class, + fetcher_class=fetcher_class, + method=method, + url_path_template=full_url_path, + base_url=base_url, + cred_lines=cred_lines, + query_block=query_block, + header_block=header_block, + description=description, + ) + return GeneratedFetcher( + module_name=module_name, + model_name=model_name, + query_params_class=qp_class, + data_class=data_class, + fetcher_class=fetcher_class, + source=source, + credentials_used=creds, + is_streaming=True, + ) + source = _render_fetcher_source( spec=spec, qp_class=qp_class, @@ -575,6 +625,155 @@ def generate_fetcher_module(spec: FetcherCommandSpec) -> GeneratedFetcher: ) +def _render_streaming_fetcher_source( + *, + spec: FetcherCommandSpec, + qp_class: str, + qp_classes: list[GeneratedClass], + data_class: str, + fetcher_class: str, + method: str, + url_path_template: str, + base_url: str, + cred_lines: list[str], + query_block: str, + header_block: str, + description: str, +) -> str: + """Format a streaming fetcher module source. + + Parameters + ---------- + spec : FetcherCommandSpec + Original command input. + qp_class : str + Query parameters class name. + qp_classes : list of GeneratedClass + Pre-flattened query parameter classes. + data_class : str + Streamed row class name. + fetcher_class : str + Fetcher class name. + method : str + HTTP verb. + url_path_template : str + URL path with placeholders intact. + base_url : str + Upstream API root, no trailing slash. + cred_lines : list of str + Credential-extraction source lines. + query_block : str + Pre-rendered query-dict source. + header_block : str + Pre-rendered headers-dict source. + description : str + First paragraph of the spec command's description. + + Returns + ------- + str + The full streaming fetcher module source. + """ + transform_query_doc = render_function_docstring( + "Validate raw input into typed query parameters.", + parameters=[("params", "dict[str, Any]", "Raw user input.")], + returns=(qp_class, "Validated query parameters."), + indent=" ", + ) + aextract_doc = render_function_docstring( + f"Open the stream at {base_url}{url_path_template}.", + parameters=[ + ("query", qp_class, "Validated query parameters."), + ("credentials", "dict | None", "Provider credentials."), + ("**kwargs", "Any", "Forwarded by the provider runtime."), + ], + returns=("AsyncIterator[str]", "Async iterator of raw stream lines."), + indent=" ", + ) + stream_data_doc = render_function_docstring( + f"Type each stream line as {data_class}.", + parameters=[ + ("query", qp_class, "Validated query parameters."), + ("data", "Any", "Async iterator of raw stream lines."), + ("**kwargs", "Any", "Forwarded by the provider runtime."), + ], + returns=(f"AsyncIterator[{data_class}]", "Async iterator of typed rows."), + indent=" ", + ) + + parts: list[str] = [] + parts.append(f'"""Streaming fetcher for {spec.name} — generated from spec."""') + parts.append("") + parts.append("from collections.abc import AsyncIterator") + parts.append("from typing import Any") + parts.append("") + parts.append("import httpx") + parts.append("") + parts.append("from openbb_core.provider.abstract.data import Data") + parts.append("from openbb_core.provider.abstract.fetcher import Fetcher") + parts.append("from openbb_core.provider.abstract.query_params import QueryParams") + parts.append("from openbb_core.provider.utils.helpers import get_querystring") + parts.append("") + parts.append("") + for cls in qp_classes: + parts.append(cls.source) + parts.append(f"class {data_class}(Data):") + parts.append(f' """Streamed row for {spec.name}."""') + parts.append("") + parts.append(" data: str") + parts.append("") + parts.append("") + parts.append( + f"class {fetcher_class}(Fetcher[{qp_class}, AsyncIterator[{data_class}]]):" + ) + parts.append(f' """Streaming fetcher for {spec.name}."""') + parts.append("") + parts.append(" @staticmethod") + parts.append(f" def transform_query(params: dict) -> {qp_class}:") + parts.append(transform_query_doc) + parts.append(f" return {qp_class}(**params)") + parts.append("") + parts.append(" @staticmethod") + parts.append( + f" async def aextract_data(query: {qp_class}, " + "credentials: dict | None, **kwargs) -> AsyncIterator[str]:" + ) + parts.append(aextract_doc) + parts.extend(cred_lines) + parts.append(query_block) + parts.append(" _query_string = get_querystring(_query_dict, [])") + parts.append( + f' _url = f"{base_url}{url_path_template}"' + ' + ("?" + _query_string if _query_string else "")' + ) + parts.append(header_block) + parts.append("") + parts.append(" async def _stream() -> AsyncIterator[str]:") + parts.append( + " async with httpx.AsyncClient(timeout=None) as _client: # noqa: S113" + ) + parts.append( + f' async with _client.stream("{method.upper()}", ' + "_url, headers=_headers) as _resp:" + ) + parts.append(" _resp.raise_for_status()") + parts.append(" async for _line in _resp.aiter_lines():") + parts.append(" if _line:") + parts.append(" yield _line") + parts.append("") + parts.append(" return _stream()") + parts.append("") + parts.append(" @staticmethod") + parts.append( + f" async def stream_data(query: {qp_class}, data: Any, **kwargs)" + f" -> AsyncIterator[{data_class}]:" + ) + parts.append(stream_data_doc) + parts.append(" async for _line in data:") + parts.append(f" yield {data_class}(data=_line)") + return "\n".join(parts).rstrip() + "\n" + + def _render_fetcher_source( *, spec: FetcherCommandSpec, diff --git a/cli/openbb_cli/codegen/router_gen.py b/cli/openbb_cli/codegen/router_gen.py index deb7ef5a31b..02ac839f05b 100644 --- a/cli/openbb_cli/codegen/router_gen.py +++ b/cli/openbb_cli/codegen/router_gen.py @@ -166,6 +166,15 @@ def _emit_router( parts.append("from openbb_core.app.query import Query") parts.append("from openbb_core.app.router import Router") + has_stream = any( + (fetcher := fetchers_by_command.get(child.full_path)) is not None + and fetcher.is_streaming + for child in node.children.values() + if child.is_command and child.cmd_spec is not None + ) + if has_stream: + parts.append("from openbb_core.app.model.stream import OBBStream") + sub_imports: list[tuple[str, str, str]] = [] for sub_name, sub_node in sorted(node.children.items()): if not sub_node.is_namespace: @@ -226,8 +235,11 @@ def _emit_router( description = (child.cmd_spec.get("description") or "").strip() fetcher = fetchers_by_command.get(child.full_path) if fetcher is not None: + renderer = ( + _render_stream_command if fetcher.is_streaming else _render_get_command + ) parts.append( - _render_get_command( + renderer( fetcher, function_name=function_name, description=description, @@ -288,6 +300,43 @@ def _render_get_command( ) +def _render_stream_command( + fetcher: GeneratedFetcher, + *, + function_name: str, + description: str, +) -> str: + """Render one ``@router.command(model=...)`` block for a streaming endpoint. + + Parameters + ---------- + fetcher : GeneratedFetcher + Per-command fetcher metadata. + function_name : str + Leaf segment of the dotted path used as the function identifier. + description : str + Spec-supplied command description. + + Returns + ------- + str + The decorator and ``async def`` block returning an OBBStream. + """ + summary = description or f"{fetcher.model_name} stream." + docstring = _format_docstring(summary) + return ( + f'@router.command(model="{fetcher.model_name}")\n' + f"async def {function_name}(\n" + " cc: CommandContext,\n" + " provider_choices: ProviderChoices,\n" + " standard_params: StandardParams,\n" + " extra_params: ExtraParams,\n" + ") -> OBBStream:\n" + f"{docstring}\n" + " return await OBBStream.from_query(Query(**locals()))\n" + ) + + def _format_docstring(text: str) -> str: """Render ``text`` as an indented function docstring. diff --git a/cli/openbb_cli/controllers/base_platform_controller.py b/cli/openbb_cli/controllers/base_platform_controller.py index ba0e0fa396c..af1a9b4ad24 100644 --- a/cli/openbb_cli/controllers/base_platform_controller.py +++ b/cli/openbb_cli/controllers/base_platform_controller.py @@ -23,6 +23,14 @@ def _OBBject() -> type: return OBBject +@lru_cache(maxsize=1) +def _OBBStream() -> type: + """Return ``OBBStream`` lazily.""" + from openbb_core.app.model.stream import OBBStream + + return OBBStream + + class DummyTranslation: """Dummy Translation for testing.""" @@ -204,7 +212,10 @@ class PlatformController(BaseController): if isinstance(obbject, list): obbject = _OBBject()(results=obbject) - if isinstance(obbject, _OBBject()): + if isinstance(obbject, _OBBStream()): + self._follow_stream(obbject) + + elif isinstance(obbject, _OBBject()): if ( session.max_obbjects_exceeded() and obbject.results @@ -312,6 +323,24 @@ class PlatformController(BaseController): ) setattr(self, f"call_{name}", bound_method) + def _follow_stream(self, stream) -> None: + """Follow a streaming (e.g. SSE) result in the foreground. + + Streams chunks to the console on a background thread and blocks until the + stream ends or the user presses Ctrl+C, then stops it cleanly. + """ + session.console.print( + f"[cyan]Streaming {stream.media_type or 'data'} — press Ctrl+C to stop.[/cyan]" + ) + stream.start() + try: + stream.wait() + except KeyboardInterrupt: + session.console.print("\n[yellow]Stream interrupted.[/yellow]") + finally: + stream.stop() + session.console.print("[cyan]Stream stopped.[/cyan]") + def _generate_controller_call(self, controller, name, parent_path, translators): """Generate controller call.""" diff --git a/cli/openbb_cli/dispatchers/http.py b/cli/openbb_cli/dispatchers/http.py index 709b37478b2..85f13493f8b 100644 --- a/cli/openbb_cli/dispatchers/http.py +++ b/cli/openbb_cli/dispatchers/http.py @@ -4,12 +4,15 @@ from __future__ import annotations import contextlib import inspect +import json import re +import warnings from datetime import datetime, timezone from time import perf_counter_ns from typing import Any import httpx +from openbb_core.app.model.abstract.warning import OpenBBWarning from openbb_cli.auth import AuthContext, AuthDecision, AuthHook from openbb_cli.dispatchers._unpack import unpack_response @@ -136,6 +139,55 @@ def _decode_response(response: httpx.Response) -> Any: return response.text +class _HttpSSESource: + """Adapt an HTTP streaming endpoint to the ``OBBStream`` body_iterator protocol. + + Holds the request details and opens the connection lazily on first iteration + (inside ``OBBStream``'s worker loop), so the stream's lifetime is bound to the + handle — not to the dispatcher's event loop. + """ + + media_type = "text/event-stream" + + def __init__(self, url, params, headers, method="get", body=None) -> None: + self.provider: str | None = None + self.warnings: list[dict[str, Any]] | None = None + self.body_iterator = self._iterate(url, params, headers, method, body) + + def _consume_headers(self, response_headers) -> None: + """Surface the server's out-of-band stream headers (provider, warnings). + + Parameters + ---------- + response_headers : Mapping + The streaming response's headers. + """ + self.provider = response_headers.get("X-OpenBB-Provider") + raw = response_headers.get("X-OpenBB-Warning") + if not raw: + return + try: + items = json.loads(raw) + except (ValueError, TypeError): + return + self.warnings = items + for item in items: + message = item.get("message") if isinstance(item, dict) else str(item) + warnings.warn(message, OpenBBWarning, stacklevel=2) + + async def _iterate(self, url, params, headers, method, body): + async with httpx.AsyncClient(timeout=None) as client: # noqa: S113 + request_kwargs: dict[str, Any] = {"params": params, "headers": headers} + if method == "post": + request_kwargs["json"] = body or {} + async with client.stream(method.upper(), url, **request_kwargs) as resp: + resp.raise_for_status() + self._consume_headers(resp.headers) + async for line in resp.aiter_lines(): + if line: + yield line + "\n" + + _NUMERIC_RESPONSE_TYPES: frozenset[str] = frozenset( {"number", "money", "percent", "double"} ) @@ -673,24 +725,47 @@ class HttpDispatcher: return override.lower() return self._command_methods.get(command, "post").lower() + def _is_streaming_command(self, command: str) -> bool: + """Return whether the command's success response is a stream.""" + cmd_spec = (self._spec_doc.get("commands") or {}).get(command) or {} + for status, content in (cmd_spec.get("response_schemas") or {}).items(): + if ( + str(status).startswith("2") + and isinstance(content, dict) + and any("event-stream" in str(ct) for ct in content) + ): + return True + return False + def _partition_params( self, command: str, params: dict[str, Any] - ) -> tuple[dict[str, Any], dict[str, Any]]: - """Split dispatch params into ``(query, body)`` by their spec ``in`` tag.""" + ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + """Split dispatch params into ``(query, body, headers)`` by spec ``in`` tag.""" cmd_spec = (self._spec_doc.get("commands") or {}).get(command) or {} - locations = { - p["name"]: p.get("in", "query") - for p in command_parameters(cmd_spec) - if isinstance(p, dict) and p.get("name") - } + locations: dict[str, tuple[str, str]] = {} + for p in command_parameters(cmd_spec): + if isinstance(p, dict) and p.get("name"): + spec_name = p["name"] + loc = p.get("in", "query") + for key in { + spec_name, + spec_name.replace("-", "_"), + spec_name.replace("_", "-"), + }: + locations[key] = (loc, spec_name) query: dict[str, Any] = {} body: dict[str, Any] = {} + headers: dict[str, Any] = {} for name, value in params.items(): - if locations.get(name, "body") == "body": + loc, spec_name = locations.get(name, ("body", name)) + if loc == "header": + if value is not None: + headers[spec_name] = str(value) + elif loc == "body": body[name] = value else: query[name] = value - return query, body + return query, body, headers async def _invoke_auth_hook( self, request: Request, method_lower: str @@ -747,14 +822,45 @@ class HttpDispatcher: message=decision.deny_reason or "auth hook denied request", ), ) - extra_headers = decision.headers or None extra_query = decision.query_params or None + query_part, body_part, header_part = self._partition_params( + wire_request.command, body_or_query + ) + extra_headers = {**(decision.headers or {}), **header_part} or None + # GET carries every non-header param in the query string; only POST + # splits the body out as JSON. + get_query = {**query_part, **body_part} + + if self._is_streaming_command(wire_request.command): + from openbb_core.app.model.stream import OBBStream + + if method_lower == "get": + merged_query = { + **self._query_params, + **(extra_query or {}), + **get_query, + } + stream_body: dict[str, Any] = {} + else: + merged_query = { + **self._query_params, + **(extra_query or {}), + **query_part, + } + stream_body = body_part + source = _HttpSSESource( + url, merged_query or None, extra_headers, method_lower, stream_body + ) + return Response( + id=request.id, ok=True, result=OBBStream(source), error=None + ) + async with self._client_context() as client: if method_lower == "get": merged_query = { **self._query_params, **(extra_query or {}), - **body_or_query, + **get_query, } if fetch_all: payload = await _fetch_all_pages( @@ -776,9 +882,6 @@ class HttpDispatcher: response.raise_for_status() payload = _decode_response(response) else: - query_part, body_part = self._partition_params( - wire_request.command, body_or_query - ) merged_query = { **self._query_params, **(extra_query or {}), diff --git a/cli/openbb_cli/dispatchers/runtime.py b/cli/openbb_cli/dispatchers/runtime.py index c6bfaf4d7b0..08d0adaddf7 100644 --- a/cli/openbb_cli/dispatchers/runtime.py +++ b/cli/openbb_cli/dispatchers/runtime.py @@ -69,10 +69,31 @@ def _to_json_line(response: Response) -> str: ) +def _looks_like_stream(value: Any) -> bool: + """Duck-type an ``OBBStream``-like handle (lazy stream with start/stop/wait).""" + return all( + callable(getattr(value, attr, None)) for attr in ("start", "stop", "wait") + ) + + +def _follow_stream(stream: Any) -> int: + """Follow a streaming result to STDOUT until it ends or Ctrl+C.""" + stream.start() + try: + stream.wait() + except KeyboardInterrupt: + pass + finally: + stream.stop() + return 0 + + def run_argv(dispatcher: Dispatcher, argv: Iterable[str]) -> int: """Dispatch a single command from argv and print the JSON response.""" request = parse_argv(argv) response = asyncio.run(_run_one(dispatcher, request)) + if response.ok and _looks_like_stream(response.result): + return _follow_stream(response.result) sys.stdout.write(_to_json_line(response) + "\n") sys.stdout.flush() return 0 if response.ok else 1 diff --git a/cli/tests/test_argparse_translator_class_processor.py b/cli/tests/test_argparse_translator_class_processor.py index 3fc35fee8fa..06ab47441a0 100644 --- a/cli/tests/test_argparse_translator_class_processor.py +++ b/cli/tests/test_argparse_translator_class_processor.py @@ -136,3 +136,27 @@ def test_custom_groups_from_reference_pulls_route_when_present(): reference={"/plain/hello": {"some": "reference-data"}}, ) rp_class.assert_called() + + +def test_custom_groups_from_reference_matches_underscore_namespace(): + """A namespace name containing ``_`` (e.g. binance_stream) resolves correctly.""" + fake_rp = MagicMock() + fake_rp.custom_groups = {"/binance_stream/book_ticker": ["grp"]} + saved = ArgparseClassProcessor._reference + ArgparseClassProcessor._reference = { + "/binance_stream/book_ticker": {"parameters": {}} + } + try: + with patch( + "openbb_cli.argparse_translator.argparse_class_processor.ReferenceToArgumentsProcessor", + return_value=fake_rp, + ) as rp_class: + out = ArgparseClassProcessor._custom_groups_from_reference( + "binance_stream", "book_ticker" + ) + rp_class.assert_called_once_with( + {"/binance_stream/book_ticker": {"parameters": {}}} + ) + assert out == ["grp"] + finally: + ArgparseClassProcessor._reference = saved diff --git a/cli/tests/test_codegen_fetcher_gen.py b/cli/tests/test_codegen_fetcher_gen.py index 15ec1ad68ee..e9b4ed99f6b 100644 --- a/cli/tests/test_codegen_fetcher_gen.py +++ b/cli/tests/test_codegen_fetcher_gen.py @@ -897,6 +897,39 @@ def test_generate_fetcher_module_get_endpoint_emits_valid_python(): assert "raise OpenBBError(" in src +def test_generate_fetcher_module_streaming_endpoint_emits_stream_data(): + spec = fg.FetcherCommandSpec( + name="crypto.feed", + cmd_spec={ + "parameters": [ + {"name": "limit", "type": "integer", "help": "Number of ticks."} + ], + "url_path": "/crypto/feed", + "method": "get", + "description": "Live crypto feed.", + "response_schemas": { + "200": {"text/event-stream": {"schema": {"type": "string"}}} + }, + "response_schema": { + "type": "object", + "properties": {"tick": {"type": "integer"}}, + }, + }, + base_url="https://api.example.com/", + api_prefix="api/v1", + provider_name="wargame", + ) + out = fg.generate_fetcher_module(spec) + + assert out.is_streaming is True + src = out.source + ast.parse(src) + assert "AsyncIterator" in src + assert "async def aextract_data" in src + assert "async def stream_data" in src + assert "httpx.AsyncClient" in src + + def test_generate_fetcher_module_path_params_substituted_into_url(): spec = fg.FetcherCommandSpec( name="us_congress.bill", diff --git a/cli/tests/test_codegen_router_gen.py b/cli/tests/test_codegen_router_gen.py index 9705a27788f..157882caa73 100644 --- a/cli/tests/test_codegen_router_gen.py +++ b/cli/tests/test_codegen_router_gen.py @@ -16,6 +16,7 @@ class _Fetcher: module_name: str fetcher_class: str model_name: str + is_streaming: bool = False @dataclass @@ -109,6 +110,33 @@ def test_generate_routers_emits_top_router_with_get_command(): ast.parse(src) +def test_generate_routers_emits_stream_command_for_streaming_fetcher(): + tree = nt.build_namespace_tree( + {"crypto.feed": {"providers": ["wargame"], "description": "Live feed"}} + ) + fetchers = { + "crypto.feed": _Fetcher( + module_name="crypto_feed", + fetcher_class="CryptoFeedFetcher", + model_name="CryptoFeed", + is_streaming=True, + ) + } + out = rg.generate_routers( + tree, + package_name="openbb_codegen", + provider_name="wargame", + fetchers_by_command=fetchers, + post_commands_by_command={}, + ) + + src = out.routers[0].source + assert "from openbb_core.app.model.stream import OBBStream" in src + assert '@router.command(model="CryptoFeed")' in src + assert "OBBStream" in src + ast.parse(src) + + def test_generate_routers_nests_sub_routers_via_include_router(): tree = nt.build_namespace_tree( { diff --git a/cli/tests/test_controllers_base_platform_controller.py b/cli/tests/test_controllers_base_platform_controller.py index a3f35e3b9c1..88191e95511 100644 --- a/cli/tests/test_controllers_base_platform_controller.py +++ b/cli/tests/test_controllers_base_platform_controller.py @@ -281,6 +281,47 @@ def test_generated_call_with_obbject_results_dispatches_to_output_adapter(mock_s mock_session.output_adapter.display.assert_called_once() +def test_generated_call_obbstream_result_is_followed(mock_session): + """A command returning an ``OBBStream`` is followed to completion.""" + from openbb_core.app.model.stream import OBBStream + + class _Source: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + yield "data: 0\n\n" + + stream = OBBStream(_Source()) + controller, _translator = _make_command_call_test_setup( + mock_session, command_returns=stream + ) + ns = MagicMock() + ns.export = "" + ns.register_obbject = False + ns.chart = False + ns.sheet_name = None + controller.parse_known_args_and_warn.return_value = ns + + controller.call_cmd([]) + assert stream.running is False + + +def test_follow_stream_handles_keyboard_interrupt(mock_session): + """``_follow_stream`` stops the stream cleanly when interrupted.""" + controller = PlatformController.__new__(PlatformController) + stream = MagicMock() + stream.media_type = "text/event-stream" + stream.wait.side_effect = KeyboardInterrupt + + controller._follow_stream(stream) + + stream.start.assert_called_once() + stream.stop.assert_called_once() + + def test_generated_call_wraps_list_in_obbject(mock_session): """A command returning a plain list gets wrapped in an ``OBBject`` before display.""" controller, translator = _make_command_call_test_setup( diff --git a/cli/tests/test_dispatchers_http.py b/cli/tests/test_dispatchers_http.py index 6ccdeeb9cbc..0e5c3a534d0 100644 --- a/cli/tests/test_dispatchers_http.py +++ b/cli/tests/test_dispatchers_http.py @@ -1118,9 +1118,217 @@ def test_partition_params_splits_query_and_body_by_in_tag(): } } d = HttpDispatcher("http://t", spec_doc=spec_doc) - query, body = d._partition_params("x.y", {"fmt": "json", "data": [1], "extra": 9}) + query, body, headers = d._partition_params( + "x.y", {"fmt": "json", "data": [1], "extra": 9} + ) assert query == {"fmt": "json"} assert body == {"data": [1], "extra": 9} + assert headers == {} + + +def test_partition_params_surfaces_header_tagged_params(): + """Params tagged ``in: header`` are routed to the headers map, not query/body.""" + spec_doc = { + "commands": { + "x.y": { + "parameters": [ + {"name": "x-api-key", "in": "header", "type": "string"}, + {"name": "fmt", "in": "query", "type": "string"}, + ], + } + } + } + d = HttpDispatcher("http://t", spec_doc=spec_doc) + query, body, headers = d._partition_params( + "x.y", {"x_api_key": "secret", "fmt": "json", "payload": 1} + ) + assert query == {"fmt": "json"} + assert body == {"payload": 1} + assert headers == {"x-api-key": "secret"} + + +def test_partition_params_drops_none_header_value(): + """A header-tagged param with a ``None`` value is omitted from headers.""" + spec_doc = { + "commands": { + "x.y": { + "parameters": [{"name": "x-api-key", "in": "header", "type": "string"}], + } + } + } + d = HttpDispatcher("http://t", spec_doc=spec_doc) + _query, _body, headers = d._partition_params("x.y", {"x_api_key": None}) + assert headers == {} + + +# --- streaming (SSE) --- + + +def test_is_streaming_command_detects_event_stream(): + spec_doc = { + "commands": { + "x.feed": {"response_schemas": {"200": {"text/event-stream": {}}}}, + "x.plain": {"response_schemas": {"200": {"application/json": {}}}}, + } + } + d = HttpDispatcher("http://t", spec_doc=spec_doc) + assert d._is_streaming_command("x.feed") is True + assert d._is_streaming_command("x.plain") is False + assert d._is_streaming_command("missing") is False + + +@pytest.mark.asyncio +async def test_dispatch_streaming_get_returns_obbstream(): + from openbb_core.app.model.stream import OBBStream + + spec_doc = { + "commands": { + "x.feed": { + "method": "get", + "response_schemas": {"200": {"text/event-stream": {}}}, + } + } + } + d = HttpDispatcher("http://t", spec_doc=spec_doc, command_methods={"x.feed": "get"}) + resp = await d.dispatch(Request(command="x.feed", params={"limit": 2})) + await d.aclose() + assert resp.ok + assert isinstance(resp.result, OBBStream) + + +@pytest.mark.asyncio +async def test_dispatch_streaming_post_returns_obbstream(): + from openbb_core.app.model.stream import OBBStream + + spec_doc = { + "commands": { + "x.feed": { + "method": "post", + "response_schemas": {"200": {"text/event-stream": {}}}, + "request_body_schema": { + "type": "object", + "properties": {"data": {"type": "array"}}, + }, + } + } + } + d = HttpDispatcher( + "http://t", spec_doc=spec_doc, command_methods={"x.feed": "post"} + ) + resp = await d.dispatch(Request(command="x.feed", params={"data": [1]})) + await d.aclose() + assert resp.ok + assert isinstance(resp.result, OBBStream) + + +@pytest.mark.asyncio +async def test_http_sse_source_iterates_get(monkeypatch): + import openbb_cli.dispatchers.http as http_mod + from openbb_cli.dispatchers.http import _HttpSSESource + + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["method"] = request.method + captured["url"] = str(request.url) + return httpx.Response(200, text="data: 0\n\ndata: 1\n") + + real_client = httpx.AsyncClient + + def make_client(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + return real_client(*args, **kwargs) + + monkeypatch.setattr(http_mod.httpx, "AsyncClient", make_client) + source = _HttpSSESource("http://t/feed", {"limit": 2}, None, "get", None) + lines = [line async for line in source.body_iterator] + assert captured["method"] == "GET" + assert "limit=2" in captured["url"] + assert any("data: 0" in line for line in lines) + + +@pytest.mark.asyncio +async def test_http_sse_source_iterates_post_with_body(monkeypatch): + import openbb_cli.dispatchers.http as http_mod + from openbb_cli.dispatchers.http import _HttpSSESource + + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["method"] = request.method + captured["body"] = request.content.decode() + return httpx.Response(200, text="data: 0\n") + + real_client = httpx.AsyncClient + + def make_client(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + return real_client(*args, **kwargs) + + monkeypatch.setattr(http_mod.httpx, "AsyncClient", make_client) + source = _HttpSSESource("http://t/feed", None, None, "post", {"data": [1]}) + lines = [line async for line in source.body_iterator] + assert captured["method"] == "POST" + assert "data" in captured["body"] + assert any("data: 0" in line for line in lines) + + +@pytest.mark.asyncio +async def test_http_sse_source_consumes_warning_header(monkeypatch): + import json as _json + + from openbb_core.app.model.abstract.warning import OpenBBWarning + + import openbb_cli.dispatchers.http as http_mod + from openbb_cli.dispatchers.http import _HttpSSESource + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + text="data: 0\n", + headers={ + "X-OpenBB-Provider": "fmp", + "X-OpenBB-Warning": _json.dumps( + [{"category": "OpenBBWarning", "message": "heads up"}] + ), + }, + ) + + real_client = httpx.AsyncClient + + def make_client(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + return real_client(*args, **kwargs) + + monkeypatch.setattr(http_mod.httpx, "AsyncClient", make_client) + source = _HttpSSESource("http://t/feed", None, None, "get", None) + with pytest.warns(OpenBBWarning, match="heads up"): + _lines = [line async for line in source.body_iterator] + assert source.provider == "fmp" + assert source.warnings is not None + assert source.warnings[0]["message"] == "heads up" + + +@pytest.mark.asyncio +async def test_http_sse_source_ignores_malformed_warning_header(monkeypatch): + import openbb_cli.dispatchers.http as http_mod + from openbb_cli.dispatchers.http import _HttpSSESource + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, text="data: 0\n", headers={"X-OpenBB-Warning": "not-json"} + ) + + real_client = httpx.AsyncClient + + def make_client(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + return real_client(*args, **kwargs) + + monkeypatch.setattr(http_mod.httpx, "AsyncClient", make_client) + source = _HttpSSESource("http://t/feed", None, None, "get", None) + _lines = [line async for line in source.body_iterator] + assert source.warnings is None # --- _help_for_provider edge: empty section --- diff --git a/cli/tests/test_dispatchers_runtime.py b/cli/tests/test_dispatchers_runtime.py index 9902c781707..bb49446166c 100644 --- a/cli/tests/test_dispatchers_runtime.py +++ b/cli/tests/test_dispatchers_runtime.py @@ -91,6 +91,53 @@ def test_run_argv_success_writes_json_and_returns_zero(capsys): assert d.requests[0].params == {"provider": "oecd"} +class _FakeStream: + """A lazy-stream stub exposing the start/wait/stop protocol.""" + + media_type = "text/event-stream" + + def __init__(self) -> None: + self.calls: list[str] = [] + + def start(self) -> None: + self.calls.append("start") + + def wait(self) -> None: + self.calls.append("wait") + + def stop(self) -> None: + self.calls.append("stop") + + +def test_run_argv_follows_stream_result(): + stream = _FakeStream() + d = _FakeDispatcher(Response(ok=True, result=stream)) + rc = run_argv(d, ["x.feed"]) + assert rc == 0 + assert stream.calls == ["start", "wait", "stop"] + + +def test_follow_stream_handles_keyboard_interrupt(): + from openbb_cli.dispatchers.runtime import _follow_stream + + class _Interrupting: + def __init__(self) -> None: + self.stopped = False + + def start(self) -> None: ... + + def wait(self) -> None: + raise KeyboardInterrupt + + def stop(self) -> None: + self.stopped = True + + stream = _Interrupting() + rc = _follow_stream(stream) + assert rc == 0 + assert stream.stopped is True + + def test_run_argv_failure_returns_one(capsys): d = _FakeDispatcher( Response(ok=False, error=ResponseError(type="X", message="nope")) diff --git a/openbb_platform/core/openbb_core/api/router/commands.py b/openbb_platform/core/openbb_core/api/router/commands.py index f2de5b78a6e..033eba63c1b 100644 --- a/openbb_platform/core/openbb_core/api/router/commands.py +++ b/openbb_platform/core/openbb_core/api/router/commands.py @@ -1,6 +1,7 @@ """Commands: generates the command map.""" import inspect +import json from collections.abc import Callable from functools import partial, wraps from inspect import Parameter, Signature, signature @@ -12,12 +13,17 @@ from fastapi.params import Depends as DependsParam from fastapi.responses import JSONResponse from fastapi.routing import APIRoute from pydantic import BaseModel +from starlette.responses import ( + Response as StarletteResponse, + StreamingResponse, +) from typing_extensions import ParamSpec from openbb_core.app.command_runner import CommandRunner from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.app.model.command_context import CommandContext from openbb_core.app.model.obbject import OBBject +from openbb_core.app.model.stream import OBBStream from openbb_core.app.model.user_settings import UserSettings from openbb_core.app.route_iter import iter_api_routes from openbb_core.app.router import RouterLoader @@ -246,7 +252,7 @@ def build_api_wrapper( @wraps(wrapped=func) async def wrapper( # noqa: PLR0912 *args: tuple[Any], **kwargs: dict[str, Any] - ) -> OBBject | JSONResponse: + ) -> OBBject | JSONResponse | StarletteResponse: user_settings: UserSettings = UserSettings.model_validate( kwargs.pop( "__authenticated_user_settings", @@ -314,6 +320,24 @@ def build_api_wrapper( output = await execute(*args, **kwargs) + if isinstance(output, StarletteResponse): + return output + if isinstance(output, OBBStream): + # The body is reserved for SSE chunks, so the response-level fields + # (id, provider, warnings) are broadcast out-of-band as headers. + stream_headers: dict[str, str] = {"X-OpenBB-Stream-Id": output.id} + if output.provider: + stream_headers["X-OpenBB-Provider"] = output.provider + if output.warnings: + stream_headers["X-OpenBB-Warning"] = json.dumps( + jsonable_encoder(output.warnings) + ) + return StreamingResponse( + output.aiter_bytes(), + media_type=output.media_type or "text/event-stream", + headers=stream_headers, + ) + if isinstance(output, OBBject): # This is where we check for `on_command_output` extensions mutated_output = getattr(output, "_extension_modified", False) diff --git a/openbb_platform/core/openbb_core/app/command_runner.py b/openbb_platform/core/openbb_core/app/command_runner.py index c8dc52380a0..55b24710210 100644 --- a/openbb_platform/core/openbb_core/app/command_runner.py +++ b/openbb_platform/core/openbb_core/app/command_runner.py @@ -19,6 +19,7 @@ from openbb_core.app.model.abstract.warning import OpenBBWarning, cast_warning from openbb_core.app.model.extension import CachedAccessor from openbb_core.app.model.metadata import Metadata from openbb_core.app.model.obbject import OBBject +from openbb_core.app.model.stream import OBBStream from openbb_core.app.provider_interface import ExtraParams from openbb_core.app.static.package_builder import PathHandler from openbb_core.env import Env @@ -247,7 +248,7 @@ class StaticCommandRunner: ) -> OBBject: """Run a command and return the output.""" obbject = await maybe_coroutine(func, **kwargs) - if isinstance(obbject, OBBject): + if isinstance(obbject, (OBBject, OBBStream)): obbject.provider = getattr( kwargs.get("provider_choices"), "provider", @@ -387,11 +388,9 @@ class StaticCommandRunner: raised_warnings = warning_list if warning_list else [] finally: if raised_warnings: - if isinstance(obbject, OBBject): - obbject.warnings = [] + if isinstance(obbject, (OBBject, OBBStream)): + obbject.warnings = [cast_warning(w) for w in raised_warnings] for w in raised_warnings: - if isinstance(obbject, OBBject): - obbject.warnings.append(cast_warning(w)) if user_settings.preferences.show_warnings: showwarning( message=w.message, @@ -447,7 +446,7 @@ class StaticCommandRunner: duration = perf_counter_ns() - start_ns if execution_context.user_settings.preferences.metadata and isinstance( - obbject, OBBject + obbject, (OBBject, OBBStream) ): try: obbject.extra["metadata"] = Metadata( @@ -470,8 +469,9 @@ class StaticCommandRunner: dep_name = to_snake_case(dep_name).replace("get_", "") dependency_param_names.add(dep_name) - for dep_key in dependency_param_names: - _ = obbject._extra_params.pop(dep_key, None) # type: ignore + if isinstance(obbject, OBBject): + for dep_key in dependency_param_names: + _ = obbject._extra_params.pop(dep_key, None) # type: ignore meta = getattr(obbject.extra.get("metadata"), "arguments", {}) diff --git a/openbb_platform/core/openbb_core/app/model/stream.py b/openbb_platform/core/openbb_core/app/model/stream.py new file mode 100644 index 00000000000..f585c69e19c --- /dev/null +++ b/openbb_platform/core/openbb_core/app/model/stream.py @@ -0,0 +1,438 @@ +"""Lazy, non-blocking, signal-safe handle for streaming responses.""" + +import asyncio +import atexit +import contextlib +import inspect +import json +import os +import signal +import sys +import threading +import time +from collections.abc import Callable +from functools import cache +from typing import TYPE_CHECKING, Any, cast + +from uuid_extensions import uuid7str + +if TYPE_CHECKING: + from openbb_core.app.query import Query + +_ACTIVE: "set[OBBStream]" = set() +_ACTIVE_LOCK = threading.Lock() + + +def stop_all(timeout: float = 2.0) -> None: + """Stop every active stream. + + Parameters + ---------- + timeout : float + Maximum seconds to wait for each stream to wind down. + """ + with _ACTIVE_LOCK: + streams = list(_ACTIVE) + for stream in streams: + with contextlib.suppress(Exception): + stream.stop(timeout=timeout) + + +@cache +def _install_shutdown_handlers() -> None: + """Install atexit and signal handlers once, chaining to existing handlers.""" + atexit.register(stop_all) + + if threading.current_thread() is not threading.main_thread(): + return + + for sig in (signal.SIGINT, getattr(signal, "SIGTERM", None)): + if sig is None: + continue + try: + previous = signal.getsignal(sig) + except (ValueError, OSError): + continue + + prev_handler = cast( + "Callable[[int, Any], Any] | None", + previous if callable(previous) else None, + ) + prev_is_default = previous == signal.SIG_DFL + + def _handler(signum, frame, _prev=prev_handler, _default=prev_is_default): + stop_all() + if _prev is not None: + _prev(signum, frame) + elif _default: + signal.signal(signum, signal.SIG_DFL) + os.kill(os.getpid(), signum) + + with contextlib.suppress(ValueError, OSError): + signal.signal(sig, _handler) + + +class OBBStream: + """Non-blocking, signal-safe handle to a streaming endpoint response. + + Iterate it directly with ``async for row in stream`` or drive it on a + background thread with ``start()`` / ``stop()``. + + Attributes + ---------- + id : str + Unique identifier of the stream. + provider : str or None + Provider name. + warnings : list or None + Warnings raised while opening the stream. + extra : dict + Extra info, e.g. execution and results metadata. + media_type : str + Stream content type, e.g. ``text/event-stream``. + + Methods + ------- + start(output=None, handler=None) + Stream on a background thread; output defaults to STDOUT and may be a + file path, file object, or socket; handler transforms or filters items. + stop(timeout=5.0) + Stop the stream and wait for it to wind down. + wait(timeout=None) + Block until the stream finishes. + + Examples + -------- + >>> stream = obb.binance_stream.book_ticker(provider="binance", symbol="BTCUSDT") + >>> stream.start() + >>> stream.stop() + >>> async for row in obb.binance_stream.book_ticker(provider="binance"): + ... ... + """ + + def __init__(self, source: Any) -> None: + """Wrap a streaming response or async iterable. + + Parameters + ---------- + source : Any + A streaming response exposing ``body_iterator`` or an async iterable. + """ + self._response = source + body_iterator = getattr(source, "body_iterator", None) + if body_iterator is None and hasattr(source, "__aiter__"): + body_iterator = source + if body_iterator is None: + raise TypeError( + "OBBStream requires a streaming response or an async iterable; " + f"got {type(source).__name__}" + ) + self._body_iterator = body_iterator + self.media_type = getattr(source, "media_type", None) or "text/event-stream" + self.id: str = uuid7str() + self.provider: str | None = None + self.warnings: list | None = None + self.extra: dict[str, Any] = {} + self._thread: threading.Thread | None = None + self._loop: asyncio.AbstractEventLoop | None = None + self._task: asyncio.Task | None = None + self._handler: Callable[[Any], Any] | None = None + self._stop_event = threading.Event() + self._started = False + + @property + def started(self) -> bool: + """Whether the stream has been started.""" + return self._started + + @property + def running(self) -> bool: + """Whether the background stream is currently running.""" + return bool(self._thread and self._thread.is_alive()) + + @classmethod + async def from_query(cls, query: "Query") -> "OBBStream": + """Build an OBBStream from a provider query. + + Parameters + ---------- + query : Query + Initialized query whose provider fetcher yields a stream. + + Returns + ------- + OBBStream + Handle over the provider's stream, carrying the query's provider and + any results metadata under ``extra``. + """ + from openbb_core.provider.abstract.annotated_result import AnnotatedResult + + result = await query.execute() + metadata = None + if isinstance(result, AnnotatedResult): + metadata = result.metadata + result = result.result + stream = cls(result) + stream.provider = getattr(query, "provider", None) + if metadata is not None: + stream.extra["results_metadata"] = metadata + return stream + + def start( + self, + output: str | os.PathLike | Any | None = None, + handler: Callable[[Any], Any] | None = None, + ) -> "OBBStream": + """Begin streaming on a background thread. + + Parameters + ---------- + output : str or os.PathLike or file or socket, optional + Destination for stream chunks. Defaults to STDOUT. + handler : Callable, optional + Function each chunk is passed through before writing. A return value + is written; ``None`` consumes the chunk. May be sync or async. + + Returns + ------- + OBBStream + This handle. + """ + if self._started: + raise RuntimeError( + "stream already started; the underlying iterator is single-use" + " and cannot be restarted — create a new stream." + ) + + _install_shutdown_handlers() + self._handler = handler + self._stop_event = threading.Event() + sink, close_sink = self._resolve_sink(output) + self._started = True + self._thread = threading.Thread( + target=self._run, args=(sink, close_sink), daemon=True + ) + with _ACTIVE_LOCK: + _ACTIVE.add(self) + self._thread.start() + return self + + def wait(self, timeout: float | None = None) -> None: + """Block until the stream finishes, processing signals in short slices. + + Parameters + ---------- + timeout : float, optional + Maximum seconds to wait. + """ + thread = self._thread + if thread is None: + return + deadline = None if timeout is None else time.monotonic() + timeout + while thread.is_alive(): + thread.join(0.2) + if deadline is not None and time.monotonic() >= deadline: + break + + def stop(self, timeout: float = 5.0) -> None: + """Stop the background stream and wait briefly for it to wind down. + + Parameters + ---------- + timeout : float + Maximum seconds to wait for the worker thread to exit. + """ + self._stop_event.set() + loop, task = self._loop, self._task + if loop is not None and task is not None and not loop.is_closed(): + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(self._cancel) + if ( + self._thread is not None + and self._thread.is_alive() + and self._thread is not threading.current_thread() + ): + self._thread.join(timeout=timeout) + with _ACTIVE_LOCK: + _ACTIVE.discard(self) + + def _cancel(self) -> None: + """Cancel the consume task; cancellation unwinds and closes the iterator.""" + if self._task is not None and not self._task.done(): + self._task.cancel() + + def __aiter__(self) -> Any: + """Return the typed async iterator of stream items.""" + self._started = True + return self._body_iterator.__aiter__() + + def aiter(self) -> Any: + """Return the underlying typed async iterator.""" + self._started = True + return self._body_iterator + + async def aiter_bytes(self) -> Any: + """Yield serialized chunks as bytes for a server StreamingResponse. + + Yields + ------ + bytes + Serialized stream chunk. + """ + self._started = True + async for chunk in self._body_iterator: + serialized = self._serialize_chunk(chunk) + yield ( + serialized + if isinstance(serialized, (bytes, bytearray)) + else serialized.encode("utf-8") + ) + + @staticmethod + def _serialize_chunk(chunk: Any) -> Any: + """Serialize a stream item to text or bytes. + + Parameters + ---------- + chunk : Any + A raw chunk, a pydantic model, or a mapping. + + Returns + ------- + str or bytes + The serialized chunk. + """ + if isinstance(chunk, (str, bytes, bytearray)): + return chunk + model_dump_json = getattr(chunk, "model_dump_json", None) + if callable(model_dump_json): + return model_dump_json() + "\n" + return json.dumps(chunk, default=str) + "\n" + + @staticmethod + def _resolve_sink(output: Any) -> tuple[Any, bool]: + """Resolve ``output`` to a writable sink and whether to close it. + + Parameters + ---------- + output : Any + A path, a writable file or pipe, a socket, or None. + + Returns + ------- + tuple + The sink object and whether this handle owns closing it. + """ + if output is None: + return sys.stdout, False + if isinstance(output, (str, os.PathLike)): + return open(output, "w", encoding="utf-8"), True # noqa: SIM115 + if hasattr(output, "write") or hasattr(output, "sendall"): + return output, False + raise TypeError( + "output must be a path, a writable file/pipe, or a socket; " + f"got {type(output).__name__}" + ) + + @classmethod + def _write(cls, sink: Any, chunk: Any) -> None: + """Write a serialized chunk to a file, STDOUT, or socket. + + Parameters + ---------- + sink : Any + The destination sink. + chunk : Any + The stream item to write. + """ + serialized = cls._serialize_chunk(chunk) + if hasattr(sink, "sendall"): + sink.sendall( + serialized.encode("utf-8") + if isinstance(serialized, str) + else serialized + ) + return + if isinstance(serialized, (bytes, bytearray)): + serialized = bytes(serialized).decode("utf-8", errors="replace") + sink.write(serialized) + if hasattr(sink, "flush"): + sink.flush() + + def _run(self, sink: Any, close_sink: bool) -> None: + """Drive the stream on a dedicated event loop in the worker thread. + + Parameters + ---------- + sink : Any + The destination sink. + close_sink : bool + Whether to close the sink when the stream ends. + """ + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + self._loop = loop + self._task = loop.create_task(self._consume(sink)) + if self._stop_event.is_set(): + self._task.cancel() + try: + with contextlib.suppress(asyncio.CancelledError, Exception): + loop.run_until_complete(self._task) + finally: + with contextlib.suppress(Exception): + self._drain_pending_tasks(loop) + with contextlib.suppress(Exception): + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + if close_sink: + with contextlib.suppress(Exception): + sink.close() + with _ACTIVE_LOCK: + _ACTIVE.discard(self) + + @staticmethod + def _drain_pending_tasks(loop: asyncio.AbstractEventLoop) -> None: + """Cancel and await tasks left on the loop, e.g. a websocket keepalive. + + Parameters + ---------- + loop : asyncio.AbstractEventLoop + The worker loop being torn down. + """ + pending = [task for task in asyncio.all_tasks(loop) if not task.done()] + if not pending: + return + for task in pending: + task.cancel() + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + + async def _consume(self, sink: Any) -> None: + """Iterate the stream, passing each chunk through the handler to the sink. + + Parameters + ---------- + sink : Any + The destination sink. + """ + handler = self._handler + iterator = self._body_iterator + try: + async for chunk in iterator: + payload = chunk + if handler is not None: + result = handler(chunk) + if inspect.isawaitable(result): + result = await result + if result is None: + continue + payload = result + self._write(sink, payload) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + with contextlib.suppress(Exception): + await aclose() + + def __repr__(self) -> str: + """Return the handle's repr.""" + return f"OBBStream(media_type={self.media_type!r}, running={self.running})" diff --git a/openbb_platform/core/openbb_core/app/router.py b/openbb_platform/core/openbb_core/app/router.py index 6c28ab5bda7..feb56c4430a 100644 --- a/openbb_platform/core/openbb_core/app/router.py +++ b/openbb_platform/core/openbb_core/app/router.py @@ -135,6 +135,11 @@ class Router: if not no_validate else func.__annotations__["return"] ) + from openbb_core.app.model.stream import OBBStream + + _ret = func.__annotations__.get("return") + if isclass(_ret) and issubclass(_ret, OBBStream): + kwargs["response_model"] = None kwargs["response_model_by_alias"] = kwargs.get( "response_model_by_alias", False ) @@ -161,6 +166,12 @@ class Router: }, ) + if isclass(_ret) and issubclass(_ret, OBBStream): + kwargs["responses"][200] = { + "description": "Event stream", + "content": {"text/event-stream": {"schema": {"type": "string"}}}, + } + # For custom deprecation if kwargs.get("deprecated", False): deprecation: OpenBBDeprecationWarning = kwargs.pop("deprecation") @@ -237,8 +248,14 @@ class SignatureInspector: cls, func: Callable[P, OBBject], model: str ) -> Callable[P, OBBject] | None: """Complete function signature.""" - if isclass(return_type := func.__annotations__["return"]) and not issubclass( - return_type, OBBject + from openbb_core.app.model.stream import OBBStream + + return_type = func.__annotations__["return"] + is_stream = isclass(return_type) and issubclass(return_type, OBBStream) + if ( + isclass(return_type) + and not issubclass(return_type, OBBject) + and not is_stream ): return func @@ -281,10 +298,11 @@ class SignatureInspector: callable_=provider_interface.params[model]["extra"], ) - func = cls.inject_return_annotation( - func=func, - annotation=provider_interface.return_annotations[model], - ) + if not is_stream: + func = cls.inject_return_annotation( + func=func, + annotation=provider_interface.return_annotations[model], + ) else: func = cls.polish_return_schema(func) diff --git a/openbb_platform/core/openbb_core/app/static/container.py b/openbb_platform/core/openbb_core/app/static/container.py index d22ff25829c..1b6b4459abe 100644 --- a/openbb_platform/core/openbb_core/app/static/container.py +++ b/openbb_platform/core/openbb_core/app/static/container.py @@ -43,6 +43,11 @@ class Container: obbject = self._command_runner.sync_run(*args, **kwargs) + if not hasattr(obbject, "_results_only") and hasattr(obbject, "body_iterator"): + from openbb_core.app.model.stream import OBBStream + + return OBBStream(obbject) + results_only = getattr(obbject, "_results_only", False) if results_only is True: diff --git a/openbb_platform/core/openbb_core/app/static/package_builder/class_definition.py b/openbb_platform/core/openbb_core/app/static/package_builder/class_definition.py index 9ba47aa4ba7..3ef3102c8e8 100644 --- a/openbb_platform/core/openbb_core/app/static/package_builder/class_definition.py +++ b/openbb_platform/core/openbb_core/app/static/package_builder/class_definition.py @@ -12,6 +12,8 @@ if TYPE_CHECKING: from importlib.util import find_spec +from openbb_core.app.model.stream import OBBStream + CHARTING_INSTALLED = find_spec("openbb_charting") is not None try: @@ -99,6 +101,14 @@ class ClassDefinition: and getattr(route, "openapi_extra", None) is not None else [] ), + response_model=( + OBBStream + if getattr( + getattr(route, "response_class", None), "media_type", None + ) + == "text/event-stream" + else getattr(route, "response_model", None) + ), ) continue diff --git a/openbb_platform/core/openbb_core/app/static/package_builder/docstring_generator.py b/openbb_platform/core/openbb_core/app/static/package_builder/docstring_generator.py index b871ca86e2b..ce2c4d10c10 100644 --- a/openbb_platform/core/openbb_core/app/static/package_builder/docstring_generator.py +++ b/openbb_platform/core/openbb_core/app/static/package_builder/docstring_generator.py @@ -212,6 +212,34 @@ class DocstringGenerator: return obbject_description + @staticmethod + def get_OBBStream_description(model_name: str) -> str: + """Get the streaming command output description.""" + indent = 2 + stream_description = ( + f"{create_indent(indent)}OBBStream\n" + f"{create_indent(indent + 1)}results : AsyncIterator[{model_name}]\n" + f"{create_indent(indent + 2)}Async iterator of streamed results;" + " use ``async for row in stream``.\n" + f"{create_indent(indent + 1)}id : str\n" + f"{create_indent(indent + 2)}Unique identifier of the stream.\n" + f"{create_indent(indent + 1)}provider : Optional[str]\n" + f"{create_indent(indent + 2)}Provider name.\n" + f"{create_indent(indent + 1)}warnings : Optional[list[Warning_]]\n" + f"{create_indent(indent + 2)}List of warnings.\n" + f"{create_indent(indent + 1)}extra : dict[str, Any]\n" + f"{create_indent(indent + 2)}Extra info.\n" + f"{create_indent(indent + 1)}start : (output=None, handler=None) -> OBBStream\n" + f"{create_indent(indent + 2)}Begin streaming on a background thread; output" + " defaults to STDOUT and may be a file path, file object, or socket;" + " handler transforms or filters each item.\n" + f"{create_indent(indent + 1)}stop : (timeout=5.0) -> None\n" + f"{create_indent(indent + 2)}Stop the stream and wait for it to wind down.\n" + f"{create_indent(indent + 1)}wait : (timeout=None) -> None\n" + f"{create_indent(indent + 2)}Block until the stream finishes.\n" + ) + return stream_description.replace("NoneType", "None") + @staticmethod def build_examples( func_path: str, @@ -252,6 +280,7 @@ class DocstringGenerator: returns: dict[str, FieldInfo], results_type: str, sections: list[str], + is_streaming: bool = False, ) -> str: """Create the docstring for model.""" docstring: str = "\n" @@ -793,8 +822,11 @@ class DocstringGenerator: docstring += "\n" docstring += f"{create_indent(2)}Returns\n" docstring += f"{create_indent(2)}-------\n" - _providers, _ = get_param_info(explicit_params.get("provider")) - docstring += cls.get_OBBject_description(results_type, _providers) + if is_streaming: + docstring += cls.get_OBBStream_description(model_name) + else: + _providers, _ = get_param_info(explicit_params.get("provider")) + docstring += cls.get_OBBject_description(results_type, _providers) # Schema underline = "-" * len(model_name) docstring += f"\n{create_indent(2)}{model_name}\n" @@ -817,6 +849,7 @@ class DocstringGenerator: formatted_params: OrderedDict[str, Parameter], model_name: str | None = None, examples: list[Example] | None = None, + return_type: Any = None, ) -> str | None: """Generate the docstring for the function.""" doc = inspect.getdoc(func) or "" @@ -831,6 +864,18 @@ class DocstringGenerator: # Map of parameter names to types param_types = {k: v.annotation for k, v in explicit_params.items()} + dependency_param_names: set = set() + try: + _sig_params = inspect.signature(func).parameters + for _pname, _p in _sig_params.items(): + _ann = _p.annotation + if isinstance(_ann, _AnnotatedAlias) and any( + hasattr(_m, "dependency") for _m in _ann.__metadata__ + ): + dependency_param_names.add(_pname) + except (AttributeError, TypeError, ValueError): + pass + if model_name: params = cls.provider_interface.params.get(model_name, {}) return_schema = cls.provider_interface.return_schema.get(model_name, None) @@ -854,6 +899,9 @@ class DocstringGenerator: and "results" in model_fields else model_name ) + is_streaming = ( + isclass(annotation) and annotation.__name__ == "OBBStream" + ) or (isclass(return_type) and return_type.__name__ == "OBBStream") doc = cls.generate_model_docstring( model_name=model_name, summary=func.__doc__ or "", @@ -862,6 +910,7 @@ class DocstringGenerator: returns=getattr(return_schema, "model_fields", {}), results_type=results_type, sections=sections, + is_streaming=is_streaming, ) doc += "\n" @@ -897,6 +946,16 @@ class DocstringGenerator: re.search(r"^\s*Examples\s*\n[-=~`]{3,}", doc, re.MULTILINE) ) + if dependency_param_names and doc_has_parameters: + doc = re.sub( + r"\n[ \t]*Parameters[ \t]*\n[-=~`]{3,}[ \t]*\n.*?" + r"(?=\n[ \t]*[A-Z][A-Za-z ]*[ \t]*\n[-=~`]{3,}|\Z)", + "\n", + doc, + flags=re.DOTALL, + ) + doc_has_parameters = False + if doc_has_returns: doc = re.sub( r"\n[ \t]*Returns[ \t]*\n[-=~`]{3,}[ \t]*\n.*?" @@ -959,8 +1018,10 @@ class DocstringGenerator: result_doc = result_doc.rstrip("\n") + "\n\n" returns_section = "Returns\n-------\n" - sig = inspect.signature(func) - return_annotation = sig.return_annotation + if return_type is not None and return_type is not inspect._empty: + return_annotation = return_type + else: + return_annotation = inspect.signature(func).return_annotation if return_annotation and return_annotation != inspect._empty: if hasattr(return_annotation, "__name__"): diff --git a/openbb_platform/core/openbb_core/app/static/package_builder/import_definition.py b/openbb_platform/core/openbb_core/app/static/package_builder/import_definition.py index 48634bc3e86..0dbb8ab4b5b 100644 --- a/openbb_platform/core/openbb_core/app/static/package_builder/import_definition.py +++ b/openbb_platform/core/openbb_core/app/static/package_builder/import_definition.py @@ -185,6 +185,15 @@ class ImportDefinition: return_type = ( sig.return_annotation if not no_validate else route.response_model or Any ) + if return_type is sig.empty and getattr(route, "response_model", None): + return_type = route.response_model + if ( + getattr(getattr(route, "response_class", None), "media_type", None) + == "text/event-stream" + ): + from openbb_core.app.model.stream import OBBStream + + return_type = OBBStream hint_type_list: list = [] diff --git a/openbb_platform/core/openbb_core/app/static/package_builder/method_definition.py b/openbb_platform/core/openbb_core/app/static/package_builder/method_definition.py index be0c66c6674..e7fabb5940a 100644 --- a/openbb_platform/core/openbb_core/app/static/package_builder/method_definition.py +++ b/openbb_platform/core/openbb_core/app/static/package_builder/method_definition.py @@ -117,6 +117,27 @@ class MethodDefinition: "fastapi.websockets.websocket", } + @staticmethod + def _is_charting_command(path: str) -> bool: + """Return whether ``path`` maps to an installed charting function. + + Parameters + ---------- + path : str + The route path. + + Returns + ------- + bool + True if charting is installed and exposes a matching function. + """ + if not CHARTING_INSTALLED: + return False + try: + return path.replace("/", "_")[1:] in Charting.functions() + except Exception: # noqa: BLE001 + return False + @staticmethod def _snake_case(name: str) -> str: if not name: @@ -235,12 +256,123 @@ class MethodDefinition: return False return True + @staticmethod + def _surfaced_param_spec(param: Parameter) -> tuple[Any, Any]: + """Resolve the ``(annotation, default)`` to surface for a dependency param. + + Unwraps ``Annotated[...]`` to the inner type, ports any description from + the FastAPI parameter object (``Header``/``Query``/``Cookie``/``Path``) + or field metadata into an ``OpenBBField``, and reads the default from the + FastAPI parameter object when present, so the surfaced parameter reads as + a plain typed, documented argument. + """ + annotation = param.annotation + inner_type = annotation + metadata: tuple = () + if isinstance(annotation, _AnnotatedAlias): + args = get_args(annotation) + inner_type = args[0] if args else Any + metadata = annotation.__metadata__ + if inner_type is Parameter.empty: + inner_type = Any + + default = param.default + + description = "" + for meta in (*metadata, default): + candidate = getattr(meta, "description", None) + if candidate: + description = candidate + break + + if default is not Parameter.empty and hasattr(default, "default"): + fastapi_default = default.default + default = ( + Parameter.empty + if fastapi_default in (Ellipsis, PydanticUndefined) + else fastapi_default + ) + + annotated = Annotated[inner_type, OpenBBField(description=description)] + return annotated, default + + @staticmethod + def _surface_dependency( + dependency_func: Callable, + ) -> tuple["OrderedDict[str, Parameter]", str] | None: + """Resolve a dependency in-process by surfacing its parameters. + + A dependency is surfaceable when every parameter can be supplied by the + caller — i.e. none is a ``Request``/``Response``/``WebSocket`` object and + none is itself a nested ``Depends(...)``. Header/Cookie/Query/plain-value + parameters are surfaced as plain arguments on the generated command and + forwarded to the dependency call, so the dependency runs without an HTTP + request (e.g. an ``x-api-key`` header becomes an ``x_api_key`` argument). + + Returns ``(surfaced_params, call_args)`` or ``None`` when the dependency + cannot be resolved in-process. + """ + try: + sig = signature(dependency_func) + except (TypeError, ValueError): + return None + + if MethodDefinition._is_none_like_return(sig.return_annotation): + return None + + surfaced: OrderedDict[str, Parameter] = OrderedDict() + arg_pairs: list[str] = [] + for pname, param in sig.parameters.items(): + annotation = param.annotation + if MethodDefinition._has_request_bound_annotation(annotation): + return None + if isinstance(annotation, _AnnotatedAlias) and any( + hasattr(meta, "dependency") for meta in annotation.__metadata__ + ): + return None + inner_type, default = MethodDefinition._surfaced_param_spec(param) + surfaced[pname] = Parameter( + name=pname, + kind=Parameter.POSITIONAL_OR_KEYWORD, + annotation=inner_type, + default=default, + ) + arg_pairs.append(f"{pname}={pname}") + return surfaced, ", ".join(arg_pairs) + + @staticmethod + def _collect_surfaced_params( + parameter_map: dict[str, Parameter], + ) -> "OrderedDict[str, Parameter]": + """Collect parameters to surface for parameter-level dependencies. + + For each ``Depends(...)`` parameter that is not already safe (resolvable + with no arguments) but is surfaceable, the dependency's own parameters are + promoted onto the generated command. + """ + surfaced: OrderedDict[str, Parameter] = OrderedDict() + for _name, param in parameter_map.items(): + if not isinstance(param.annotation, _AnnotatedAlias): + continue + for meta in param.annotation.__metadata__: + dependency_func = getattr(meta, "dependency", None) + if dependency_func is None: + continue + if MethodDefinition._is_safe_dependency(dependency_func): + continue + result = MethodDefinition._surface_dependency(dependency_func) + if result is not None: + params, _ = result + for pname, pparam in params.items(): + surfaced.setdefault(pname, pparam) + return surfaced + @staticmethod def build_class_loader_method(path: str) -> str: """Build the class loader method.""" module_name = PathHandler.build_module_name(path=path) class_name = PathHandler.build_module_class(path=path) - function_name = path.rsplit("/", maxsplit=1)[-1].strip("/") + function_name = path.rsplit("/", maxsplit=1)[-1].strip("/").replace("-", "_") description = PathHandler.get_router_description(path) code = "\n @property\n" @@ -536,7 +668,7 @@ class MethodDefinition: path_params = PathHandler.extract_path_parameters(path) # we need to add the chart parameter here bc of the docstring generation - if CHARTING_INSTALLED and path.replace("/", "_")[1:] in Charting.functions(): + if MethodDefinition._is_charting_command(path): parameter_map["chart"] = Parameter( name="chart", kind=Parameter.POSITIONAL_OR_KEYWORD, @@ -1112,6 +1244,7 @@ class MethodDefinition: formatted_params: OrderedDict[str, Parameter], model_name: str | None = None, examples: list[Example] | None = None, + return_type: Any = None, ): """Build the command method docstring.""" doc = func.__doc__ @@ -1121,6 +1254,7 @@ class MethodDefinition: formatted_params=formatted_params, model_name=model_name, examples=examples, + return_type=return_type, ) code = ( @@ -1169,11 +1303,18 @@ class MethodDefinition: if not (hasattr(meta, "dependency") and meta.dependency is not None): continue dependency_func = meta.dependency - if not MethodDefinition._is_safe_dependency(dependency_func): - continue func_name = dependency_func.__name__ - dependency_calls.append(f" {name} = {func_name}()") - dependency_names.add(name) + if MethodDefinition._is_safe_dependency(dependency_func): + dependency_calls.append(f" {name} = {func_name}()") + dependency_names.add(name) + elif ( + surfaced := MethodDefinition._surface_dependency(dependency_func) + ) is not None: + _, call_args = surfaced + dependency_calls.append( + f" {name} = {func_name}({call_args})" + ) + dependency_names.add(name) return dependency_calls, dependency_names @@ -1200,7 +1341,7 @@ class MethodDefinition: if dependency_calls: code += "\n".join(dependency_calls) + "\n\n" - if CHARTING_INSTALLED and path.replace("/", "_")[1:] in Charting.functions(): + if MethodDefinition._is_charting_command(path): parameter_map["chart"] = Parameter( name="chart", kind=Parameter.POSITIONAL_OR_KEYWORD, @@ -1226,6 +1367,14 @@ class MethodDefinition: has_extra_params = False for name, param in parameter_map.items(): + is_dependency = isinstance(param.annotation, _AnnotatedAlias) and any( + getattr(meta, "dependency", None) is not None + for meta in param.annotation.__metadata__ + ) + if is_dependency: + if name in dependency_names: + code += f" {name}={name},\n" + continue if name == "extra_params": has_extra_params = True fields = ( @@ -1276,24 +1425,17 @@ class MethodDefinition: ) and not MethodDefinition.is_data_processing_function(path) ): - has_depends = any( - hasattr(meta, "dependency") - for meta in param.annotation.__metadata__ + model = param.annotation.__args__[0] + fields = getattr( + model, + "model_fields", + getattr(model, "__pydantic_fields__", {}), ) - if not has_depends: - model = param.annotation.__args__[0] - fields = getattr( - model, - "model_fields", - getattr(model, "__pydantic_fields__", {}), - ) - values = {k: k for k in fields} - code += f" {name}={{\n" - for k, v in values.items(): - code += f' "{k}": {v},\n' - code += " },\n" - else: - code += f" {name}={name},\n" + values = {k: k for k in fields} + code += f" {name}={{\n" + for k, v in values.items(): + code += f' "{k}": {v},\n' + code += " },\n" elif isclass(param.annotation) and issubclass( param.annotation, QueryParams ): @@ -1357,11 +1499,12 @@ class MethodDefinition: func: Callable, model_name: str | None = None, examples: list[Example] | None = None, + response_model: Any = None, ) -> str: """Build the command method.""" path_parts = [p for p in path.split("/") if p and not p.startswith("{")] func_name = ( - path_parts[-1] if path_parts else func.__name__ # ty: ignore[unresolved-attribute] + path_parts[-1].replace("-", "_") if path_parts else func.__name__ # ty: ignore[unresolved-attribute] ) sig = signature(func) parameter_map = dict(sig.parameters) @@ -1434,6 +1577,10 @@ class MethodDefinition: if name not in parameter_map: parameter_map[name] = param + for name, param in cls._collect_surfaced_params(parameter_map).items(): + if name not in parameter_map: + parameter_map[name] = param + formatted_params = cls.format_params( path=path, parameter_map=parameter_map, func=func ) @@ -1451,10 +1598,14 @@ class MethodDefinition: default=Parameter.empty, ) + return_type = sig.return_annotation + if return_type is signature(func).empty and response_model is not None: + return_type = response_model + code = cls.build_command_method_signature( func_name=func_name, formatted_params=formatted_params, - return_type=sig.return_annotation, + return_type=return_type, path=path, model_name=model_name, ) @@ -1464,6 +1615,7 @@ class MethodDefinition: formatted_params=formatted_params, model_name=model_name, examples=examples, + return_type=return_type, ) code += cls.build_command_method_body( diff --git a/openbb_platform/core/openbb_core/provider/abstract/fetcher.py b/openbb_platform/core/openbb_core/provider/abstract/fetcher.py index 6004c379618..c876c2dbfc6 100644 --- a/openbb_platform/core/openbb_core/provider/abstract/fetcher.py +++ b/openbb_platform/core/openbb_core/provider/abstract/fetcher.py @@ -56,6 +56,27 @@ class Fetcher(Generic[Q, R]): """Transform the provider-specific data.""" raise NotImplementedError + @staticmethod + def stream_data(query: Q, data: Any, **kwargs) -> R: + """Transform a stream of extracted items into typed data. + + Override with an ``async def`` generator for streaming fetchers; ``data`` + is the async iterator returned by ``extract_data``. + + Parameters + ---------- + query : Q + The validated query. + data : Any + Async iterator of raw items from ``extract_data``. + + Returns + ------- + R + Async iterator of typed data rows. + """ + raise NotImplementedError + def __init_subclass__(cls, *args, **kwargs): """Initialize the subclass.""" super().__init_subclass__(*args, **kwargs) @@ -81,8 +102,15 @@ class Fetcher(Generic[Q, R]): data = await maybe_coroutine( cls.extract_data, query=query, credentials=credentials, **kwargs ) + if cls.is_streaming: + return cls.stream_data(query=query, data=data, **kwargs) return cls.transform_data(query=query, data=data, **kwargs) + @classproperty + def is_streaming(self) -> bool: + """Whether this fetcher streams its results via ``stream_data``.""" + return self.stream_data is not Fetcher.stream_data + @classproperty def query_params_type(self) -> Q: """Get the type of query.""" @@ -103,9 +131,38 @@ class Fetcher(Generic[Q, R]): @staticmethod def _get_data_type(data: Any) -> D: - """Get the type of the data.""" - if get_origin(data) is list: - data = get_args(data)[0] + """Get the type of the data. + + Parameters + ---------- + data : Any + The fetcher return annotation, possibly a container of the row type. + + Returns + ------- + D + The row data type, unwrapped from list or async-iterator containers. + """ + from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Iterable, + Iterator, + ) + + containers = ( + list, + AsyncIterator, + AsyncIterable, + AsyncGenerator, + Iterator, + Iterable, + ) + if get_origin(data) in containers: + args = get_args(data) + if args: + data = args[0] return data @classmethod @@ -136,11 +193,6 @@ class Fetcher(Generic[Q, R]): DataFrame = require_optional("pandas").DataFrame # type: ignore[union-attr] query = cls.transform_query(params=params) - data = run_async( - cls.extract_data, query=query, credentials=credentials, **kwargs - ) - result = cls.transform_data(query=query, data=data, **kwargs) - data_type_fields = cls.data_type.model_fields # Class Assertions assert isinstance(cls.require_credentials, bool), ( @@ -156,6 +208,16 @@ class Fetcher(Generic[Q, R]): f"Query must have the correct values. Expected: {params} Got: {query.__dict__}" ) + if cls.is_streaming: + cls._test_stream(query, credentials, **kwargs) + return + + data = run_async( + cls.extract_data, query=query, credentials=credentials, **kwargs + ) + result = cls.transform_data(query=query, data=data, **kwargs) + data_type_fields = cls.data_type.model_fields + # Data Assertions if not isinstance(data, DataFrame): assert data, "Data must not be None." @@ -241,3 +303,66 @@ class Fetcher(Generic[Q, R]): assert issubclass(type(transformed_data), cls.return_type), ( f"Transformed data must be of the correct type. Expected: {cls.return_type} Got: {type(transformed_data)}" ) + + @classmethod + def _test_stream( + cls, + query: Q, + credentials: dict[str, str] | None = None, + max_items: int = 1, + **kwargs, + ) -> None: + """Test a streaming fetcher by consuming and validating stream items. + + Parameters + ---------- + query : Q + The validated query. + credentials : dict[str, str] or None + The credentials to test the fetcher with. + max_items : int + Number of items to consume before stopping. + + Raises + ------ + AssertionError + If any of the streaming tests fail. + """ + data_type = cls.data_type + data_type_fields = data_type.model_fields + + async def _run() -> None: + from collections.abc import AsyncIterator + from typing import cast + + raw = await maybe_coroutine( + cls.extract_data, query=query, credentials=credentials, **kwargs + ) + assert hasattr(raw, "__aiter__"), ( + "extract_data must return an async iterable for a streaming fetcher." + ) + stream = cls.stream_data(query=query, data=raw, **kwargs) + assert hasattr(stream, "__aiter__"), ( + "stream_data must return an async iterable for a streaming fetcher." + ) + iterator = cast("AsyncIterator[Any]", stream).__aiter__() + count = 0 + try: + while count < max_items: + item = await iterator.__anext__() + assert issubclass(type(item), data_type), ( + f"Streamed item type mismatch. Expected: {data_type} Got: {type(item)}" + ) + assert all(field in item.__dict__ for field in data_type_fields), ( + f"Streamed item must have the correct fields. Expected: {data_type_fields} Got: {item.__dict__}" + ) + count += 1 + except StopAsyncIteration: + pass + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + assert count > 0, "Stream must yield at least one item." + + run_async(_run) diff --git a/openbb_platform/core/openbb_core/provider/registry_map.py b/openbb_platform/core/openbb_core/provider/registry_map.py index f9184f3a382..6efbd342ebb 100644 --- a/openbb_platform/core/openbb_core/provider/registry_map.py +++ b/openbb_platform/core/openbb_core/provider/registry_map.py @@ -143,8 +143,30 @@ class RegistryMap: @staticmethod def _get_results_type(fetcher: Fetcher) -> Any: - """Extract return info from fetcher.""" - return get_origin(getattr(fetcher, "return_type", None)) + """Extract return info from fetcher. + + Parameters + ---------- + fetcher : Fetcher + The fetcher whose return container is inspected. + + Returns + ------- + Any + The results container origin; streaming containers normalize to list. + """ + from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Iterable, + Iterator, + ) + + origin = get_origin(getattr(fetcher, "return_type", None)) + if origin in (AsyncIterator, AsyncIterable, AsyncGenerator, Iterator, Iterable): + return list + return origin @staticmethod def _extract_info( diff --git a/openbb_platform/core/tests/api/test_router/test_commands_helpers.py b/openbb_platform/core/tests/api/test_router/test_commands_helpers.py index c40220c361a..fe5f3bf8c05 100644 --- a/openbb_platform/core/tests/api/test_router/test_commands_helpers.py +++ b/openbb_platform/core/tests/api/test_router/test_commands_helpers.py @@ -225,6 +225,103 @@ async def test_wrapper_passes_through_non_obbject(): assert out == {"plain": "dict"} +@pytest.mark.asyncio +async def test_wrapper_returns_streaming_response_for_obbstream(): + from starlette.responses import StreamingResponse + + from openbb_core.app.model.stream import OBBStream + + runner = MagicMock(spec=CommandRunner) + + class _Source: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + yield "data: 0\n\n" + + async def _run(path, user_settings, *args, **kwargs): + return OBBStream(_Source()) + + runner.run = _run + + async def endpoint(symbol: str = "AAPL", **kwargs): + return OBBject(results=[]) + + route = _make_route(endpoint) + wrapper = build_api_wrapper(runner, route) + out = await wrapper(symbol="AAPL") + assert isinstance(out, StreamingResponse) + assert out.media_type == "text/event-stream" + assert out.headers["X-OpenBB-Stream-Id"] + + +@pytest.mark.asyncio +async def test_wrapper_streaming_broadcasts_warnings_and_provider_headers(): + import json as _json + + from starlette.responses import StreamingResponse + + from openbb_core.app.model.abstract.warning import Warning_ + from openbb_core.app.model.stream import OBBStream + + runner = MagicMock(spec=CommandRunner) + + class _Source: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + yield "data: 0\n\n" + + stream = OBBStream(_Source()) + stream.provider = "fmp" + stream.warnings = [Warning_(category="OpenBBWarning", message="heads up")] + + async def _run(path, user_settings, *args, **kwargs): + return stream + + runner.run = _run + + async def endpoint(symbol: str = "AAPL", **kwargs): + return OBBject(results=[]) + + route = _make_route(endpoint) + wrapper = build_api_wrapper(runner, route) + out = await wrapper(symbol="AAPL") + assert isinstance(out, StreamingResponse) + assert out.headers["X-OpenBB-Provider"] == "fmp" + assert out.headers["X-OpenBB-Stream-Id"] == stream.id + broadcast = _json.loads(out.headers["X-OpenBB-Warning"]) + assert broadcast[0]["category"] == "OpenBBWarning" + assert broadcast[0]["message"] == "heads up" + + +@pytest.mark.asyncio +async def test_wrapper_passes_through_response(): + from starlette.responses import PlainTextResponse + + runner = MagicMock(spec=CommandRunner) + response = PlainTextResponse("hi") + + async def _run(path, user_settings, *args, **kwargs): + return response + + runner.run = _run + + async def endpoint(symbol: str = "AAPL", **kwargs): + return OBBject(results=[]) + + route = _make_route(endpoint) + wrapper = build_api_wrapper(runner, route) + out = await wrapper(symbol="AAPL") + assert out is response + + @pytest.mark.asyncio async def test_wrapper_no_validate_route_returns_jsonresponse(): runner = MagicMock(spec=CommandRunner) diff --git a/openbb_platform/core/tests/app/model/test_stream.py b/openbb_platform/core/tests/app/model/test_stream.py new file mode 100644 index 00000000000..c0ffb02c1d5 --- /dev/null +++ b/openbb_platform/core/tests/app/model/test_stream.py @@ -0,0 +1,425 @@ +"""Tests for the OBBStream model.""" + +import asyncio +import io +import socket +import time + +import pytest + +from openbb_core.app.model import stream as stream_module +from openbb_core.app.model.stream import ( + OBBStream, + _install_shutdown_handlers, + stop_all, +) +from openbb_core.provider.abstract.annotated_result import AnnotatedResult + + +class _Source: + """A streaming response stub exposing ``body_iterator``.""" + + media_type = "text/event-stream" + + def __init__(self, count: int = 3, sleep: float = 0.0): + self._count = count + self._sleep = sleep + self.body_iterator = self._gen() + + async def _gen(self): + for index in range(self._count): + if self._sleep: + await asyncio.sleep(self._sleep) + yield f"data: {index}\n\n" + + +class _Model: + """A row stub exposing ``model_dump_json``.""" + + def model_dump_json(self) -> str: + return '{"tick": 1}' + + +def test_init_with_body_iterator(): + handle = OBBStream(_Source()) + assert handle.media_type == "text/event-stream" + assert handle.started is False + assert handle.running is False + assert "OBBStream(" in repr(handle) + + +def test_init_with_async_iterable(): + async def _agen(): + yield "x" + + handle = OBBStream(_agen()) + assert handle.media_type == "text/event-stream" + + +def test_init_rejects_non_iterable(): + with pytest.raises(TypeError): + OBBStream(object()) + + +def test_resolve_sink_rejects_invalid(): + with pytest.raises(TypeError): + OBBStream._resolve_sink(123) + + +def test_serialize_chunk_variants(): + assert OBBStream._serialize_chunk("a") == "a" + assert OBBStream._serialize_chunk(b"b") == b"b" + assert OBBStream._serialize_chunk(_Model()) == '{"tick": 1}\n' + assert OBBStream._serialize_chunk({"x": 1}) == '{"x": 1}\n' + + +def test_start_to_file(tmp_path): + path = tmp_path / "out.txt" + handle = OBBStream(_Source(3)) + handle.start(output=str(path)) + handle.wait(timeout=2) + handle.stop() + assert "data: 0" in path.read_text() + assert handle.running is False + + +def test_start_to_stringio_default_repr(): + buf = io.StringIO() + handle = OBBStream(_Source(3)) + handle.start(output=buf) + handle.wait(timeout=2) + handle.stop() + assert handle.started is True + lines = [line for line in buf.getvalue().splitlines() if line.strip()] + assert lines == ["data: 0", "data: 1", "data: 2"] + + +def test_start_to_socket(): + recv, send = socket.socketpair() + recv.settimeout(2) + handle = OBBStream(_Source(2)) + handle.start(output=send) + handle.wait(timeout=2) + handle.stop() + data = recv.recv(4096) + recv.close() + send.close() + assert b"data: 0" in data + + +def test_start_serializes_model_rows(): + buf = io.StringIO() + + class _ModelSource: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + yield _Model() + + handle = OBBStream(_ModelSource()) + handle.start(output=buf) + handle.wait(timeout=2) + handle.stop() + assert buf.getvalue() == '{"tick": 1}\n' + + +def test_start_twice_raises(): + handle = OBBStream(_Source(1000, sleep=0.05)) + handle.start(output=io.StringIO()) + try: + with pytest.raises(RuntimeError): + handle.start() + finally: + handle.stop() + + +def test_start_after_stop_raises(): + handle = OBBStream(_Source(3)) + handle.start(output=io.StringIO()) + handle.wait(timeout=2) + handle.stop() + with pytest.raises(RuntimeError, match="single-use"): + handle.start() + + +def test_start_after_iteration_raises(): + async def _run(): + handle = OBBStream(_Source(2)) + _ = [chunk async for chunk in handle] + return handle + + handle = asyncio.run(_run()) + with pytest.raises(RuntimeError, match="single-use"): + handle.start() + + +def test_handler_filters_chunks(): + buf = io.StringIO() + seen: list = [] + + def _handler(chunk): + seen.append(chunk) + + handle = OBBStream(_Source(3)) + handle.start(output=buf, handler=_handler) + handle.wait(timeout=2) + handle.stop() + assert len(seen) == 3 + assert buf.getvalue() == "" + + +def test_handler_async_transform(): + buf = io.StringIO() + + async def _handler(chunk): + return chunk.upper() + + handle = OBBStream(_Source(2)) + handle.start(output=buf, handler=_handler) + handle.wait(timeout=2) + handle.stop() + assert "DATA: 0" in buf.getvalue() + + +def test_async_iteration(): + async def _run(): + return [chunk async for chunk in OBBStream(_Source(3))] + + assert asyncio.run(_run()) == ["data: 0\n\n", "data: 1\n\n", "data: 2\n\n"] + + +def test_aiter_bytes(): + async def _run(): + return [chunk async for chunk in OBBStream(_Source(2)).aiter_bytes()] + + assert asyncio.run(_run()) == [b"data: 0\n\n", b"data: 1\n\n"] + + +def test_aiter_returns_iterator(): + handle = OBBStream(_Source(1)) + assert handle.aiter() is handle._body_iterator + assert handle.started is True + + +def test_from_query_plain_result(): + class _Query: + async def execute(self): + async def _gen(): + yield "a" + + return _gen() + + async def _run(): + handle = await OBBStream.from_query(_Query()) + return [chunk async for chunk in handle] + + assert asyncio.run(_run()) == ["a"] + + +def test_from_query_annotated_result(): + class _Query: + async def execute(self): + async def _gen(): + yield "b" + + return AnnotatedResult(result=_gen(), metadata={}) + + async def _run(): + handle = await OBBStream.from_query(_Query()) + return [chunk async for chunk in handle] + + assert asyncio.run(_run()) == ["b"] + + +def test_wait_returns_immediately_when_not_started(): + handle = OBBStream(_Source(1)) + handle.wait(timeout=1) + assert handle.running is False + + +def test_stop_interrupts_stalled_stream(): + class _Stalled: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + await asyncio.sleep(100) + yield "never" + + handle = OBBStream(_Stalled()) + handle.start(output=io.StringIO()) + time.sleep(0.1) + started = time.monotonic() + handle.stop(timeout=5) + assert time.monotonic() - started < 3 + assert handle.running is False + + +def test_run_cancels_when_stop_requested_before_loop_starts(): + handle = OBBStream(_Source(3)) + handle._stop_event.set() + buf = io.StringIO() + try: + handle._run(buf, False) + finally: + asyncio.set_event_loop(None) + assert handle.running is False + assert buf.getvalue() == "" + + +def test_drain_pending_tasks_cancels_leftover_loop_tasks(): + class _LeakySource: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + asyncio.ensure_future(asyncio.sleep(1000)) + yield "data: 0\n\n" + + handle = OBBStream(_LeakySource()) + handle.start(output=io.StringIO()) + handle.wait(timeout=2) + handle.stop() + assert handle.running is False + + +def test_stop_when_never_started_is_noop(): + handle = OBBStream(_Source(1)) + handle.stop() + assert handle.running is False + + +def test_wait_with_timeout_on_long_stream(): + handle = OBBStream(_Source(1000, sleep=0.05)) + handle.start(output=io.StringIO()) + handle.wait(timeout=0.2) + handle.stop() + assert handle.running is False + + +def test_stop_all_stops_active_streams(): + handle = OBBStream(_Source(1000, sleep=0.05)) + handle.start(output=io.StringIO()) + assert handle.running is True + stop_all() + assert handle.running is False + + +def test_install_shutdown_handlers_idempotent_and_signal_chain(monkeypatch): + import signal + + _install_shutdown_handlers.cache_clear() + + registered: dict = {} + + def _fake_signal(sig, handler): + registered[sig] = handler + + monkeypatch.setattr(signal, "signal", _fake_signal) + monkeypatch.setattr(signal, "getsignal", lambda _sig: signal.SIG_DFL) + monkeypatch.setattr(stream_module, "stop_all", lambda *a, **k: None) + + _install_shutdown_handlers() + _install_shutdown_handlers() + + handler = registered.get(signal.SIGINT) + assert callable(handler) + + killed: list = [] + monkeypatch.setattr(stream_module.os, "kill", lambda pid, sig: killed.append(sig)) + handler(signal.SIGINT, None) + assert killed == [signal.SIGINT] + + _install_shutdown_handlers.cache_clear() + + +def test_resolve_sink_defaults_to_stdout(): + import sys + + sink, close = OBBStream._resolve_sink(None) + assert sink is sys.stdout + assert close is False + + +def test_write_decodes_bytes_for_text_sink(): + buf = io.StringIO() + OBBStream._write(buf, b"hello\n") + assert buf.getvalue() == "hello\n" + + +def test_install_handlers_non_main_thread(): + import threading + + _install_shutdown_handlers.cache_clear() + errors: list = [] + + def _run(): + try: + _install_shutdown_handlers() + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + worker = threading.Thread(target=_run) + worker.start() + worker.join() + assert not errors + _install_shutdown_handlers.cache_clear() + + +def test_install_handlers_getsignal_raises(monkeypatch): + import signal + + _install_shutdown_handlers.cache_clear() + + def _raise(_sig): + raise ValueError("no handler") + + monkeypatch.setattr(signal, "signal", lambda *a: None) + monkeypatch.setattr(signal, "getsignal", _raise) + monkeypatch.setattr(stream_module, "stop_all", lambda *a, **k: None) + _install_shutdown_handlers() + _install_shutdown_handlers.cache_clear() + + +def test_install_handlers_sigterm_absent(monkeypatch): + import signal + + _install_shutdown_handlers.cache_clear() + monkeypatch.delattr(signal, "SIGTERM", raising=False) + monkeypatch.setattr(signal, "signal", lambda *a: None) + monkeypatch.setattr(signal, "getsignal", lambda _sig: signal.SIG_DFL) + monkeypatch.setattr(stream_module, "stop_all", lambda *a, **k: None) + _install_shutdown_handlers() + _install_shutdown_handlers.cache_clear() + + +def test_install_shutdown_handlers_chains_previous(monkeypatch): + import signal + + _install_shutdown_handlers.cache_clear() + + registered: dict = {} + prev_calls: list = [] + + def _prev(signum, frame): + prev_calls.append(signum) + + def _record_signal(sig, handler): + registered[sig] = handler + + monkeypatch.setattr(signal, "signal", _record_signal) + monkeypatch.setattr(signal, "getsignal", lambda _sig: _prev) + monkeypatch.setattr(stream_module, "stop_all", lambda *a, **k: None) + + _install_shutdown_handlers() + registered[signal.SIGINT](signal.SIGINT, None) + assert prev_calls == [signal.SIGINT] + + _install_shutdown_handlers.cache_clear() diff --git a/openbb_platform/core/tests/app/static/package_builder/test_docstring_generator_helpers.py b/openbb_platform/core/tests/app/static/package_builder/test_docstring_generator_helpers.py index e2cca4d4afd..c12b5288ed8 100644 --- a/openbb_platform/core/tests/app/static/package_builder/test_docstring_generator_helpers.py +++ b/openbb_platform/core/tests/app/static/package_builder/test_docstring_generator_helpers.py @@ -260,6 +260,36 @@ def test_get_OBBject_description_default_providers(): assert "Optional[str]" in out +def test_get_OBBStream_description_mirrors_obbject_field_format(): + out = DocstringGenerator.get_OBBStream_description("MyFeed") + assert "OBBStream" in out + assert "results : AsyncIterator[MyFeed]" in out + assert "id : str" in out + assert "provider : Optional[str]" in out + assert "warnings : Optional[list[Warning_]]" in out + assert "extra : dict[str, Any]" in out + assert "start : (output=None, handler=None) -> OBBStream" in out + assert "stop : (timeout=5.0) -> None" in out + assert "wait : (timeout=None) -> None" in out + + +def test_generate_model_docstring_streaming_emits_obbstream_returns(): + out = DocstringGenerator.generate_model_docstring( + model_name="MyFeed", + summary="Stream it.", + explicit_params={}, + kwarg_params={}, + returns={}, + results_type="MyFeed", + sections=["returns"], + is_streaming=True, + ) + assert "OBBStream" in out + assert "results : AsyncIterator[MyFeed]" in out + # The OBBject result block must not be used for a stream. + assert "Serializable results." not in out + + def test_get_field_type_handles_forward_ref_extended(): from typing import ForwardRef diff --git a/openbb_platform/core/tests/app/static/package_builder/test_method_definition_helpers.py b/openbb_platform/core/tests/app/static/package_builder/test_method_definition_helpers.py index 6651eb3f4de..7e1b0bbafa6 100644 --- a/openbb_platform/core/tests/app/static/package_builder/test_method_definition_helpers.py +++ b/openbb_platform/core/tests/app/static/package_builder/test_method_definition_helpers.py @@ -1467,3 +1467,187 @@ def test_build_command_method_getsource_typeerror(monkeypatch): monkeypatch.setattr(md.inspect, "getsource", _patched_getsource) code = MethodDefinition.build_command_method("/x/y", endpoint) assert "def y" in code + + +def test_surface_dependency_none_like_return(): + def dep() -> None: ... + + assert MethodDefinition._surface_dependency(dep) is None + + +def test_surface_dependency_request_bound_param(): + def dep(request: Request) -> dict: ... + + assert MethodDefinition._surface_dependency(dep) is None + + +def test_surface_dependency_nested_depends(): + from fastapi import Depends + + def other() -> int: + return 1 + + def dep(value: Annotated[int, Depends(other)]) -> dict: ... + + assert MethodDefinition._surface_dependency(dep) is None + + +def test_surface_dependency_surfaceable_header(): + from fastapi import Header + + def dep(x_api_key: Annotated[str, Header()]) -> dict: ... + + result = MethodDefinition._surface_dependency(dep) + assert result is not None + surfaced, call_args = result + assert "x_api_key" in surfaced + assert call_args == "x_api_key=x_api_key" + + +def test_surfaced_param_spec_empty_annotation_defaults_to_any(): + from typing import Any, get_args + + param = Parameter("x", Parameter.POSITIONAL_OR_KEYWORD) + annotation, default = MethodDefinition._surfaced_param_spec(param) + assert get_args(annotation)[0] is Any + assert default is Parameter.empty + + +def test_surfaced_param_spec_reads_fastapi_default(): + from fastapi import Query as FastAPIQuery + + param = Parameter( + "x", + Parameter.POSITIONAL_OR_KEYWORD, + default=FastAPIQuery(default=5), + annotation=int, + ) + _annotation, default = MethodDefinition._surfaced_param_spec(param) + assert default == 5 + + +def test_collect_surfaced_params_promotes_header(): + from fastapi import Depends, Header + + def dep(x_api_key: Annotated[str, Header()]) -> dict: ... + + parameter_map = { + "principal": Parameter( + "principal", + Parameter.POSITIONAL_OR_KEYWORD, + annotation=Annotated[dict, Depends(dep)], + ) + } + surfaced = MethodDefinition._collect_surfaced_params(parameter_map) + assert "x_api_key" in surfaced + + +def test_is_charting_command_without_charting(): + with patch.object(md_module, "CHARTING_INSTALLED", False): + assert MethodDefinition._is_charting_command("/x/y") is False + + +def test_is_charting_command_handles_exception(): + with ( + patch.object(md_module, "CHARTING_INSTALLED", True), + patch.object(md_module, "Charting", create=True) as mock_charting, + ): + mock_charting.functions.side_effect = RuntimeError("boom") + assert MethodDefinition._is_charting_command("/x/y") is False + + +def test_format_params_flattens_query_params_model(): + from openbb_core.provider.abstract.query_params import QueryParams + + class _MyQueryParams(QueryParams): + symbol: str + limit: int = 10 + + parameter_map = { + "params": Parameter( + "params", + Parameter.POSITIONAL_OR_KEYWORD, + annotation=Annotated[_MyQueryParams, "meta"], + ) + } + out = MethodDefinition.format_params("/x/y", parameter_map) + assert "symbol" in out + assert "limit" in out + + +def test_build_command_method_body_query_params_model(monkeypatch): + from openbb_core.provider.abstract.query_params import QueryParams + + class _BodyQueryParams(QueryParams): + symbol: str = "AAPL" + + monkeypatch.setattr( + MethodDefinition, "is_deprecated_function", staticmethod(lambda p: False) + ) + + def endpoint(params: _BodyQueryParams): + return params + + code = MethodDefinition.build_command_method_body("/x/y", endpoint) + assert '"symbol": symbol' in code + + +def test_build_command_method_body_pydantic_model_param(monkeypatch): + from pydantic import BaseModel + + class _Body(BaseModel): + ticker: str = "AAPL" + + monkeypatch.setattr( + MethodDefinition, "is_deprecated_function", staticmethod(lambda p: False) + ) + + def endpoint(payload: Annotated[_Body, OpenBBField(description="body")]): + return payload + + code = MethodDefinition.build_command_method_body("/x/y", endpoint) + assert '"ticker": ticker' in code + + +def test_format_params_query_params_model_data_processing_path(monkeypatch): + from openbb_core.provider.abstract.query_params import QueryParams + + class _MyQueryParams(QueryParams): + data: str + start_date: str = "2020-01-01" + tags: list = Field(default_factory=list) + + monkeypatch.setattr( + MethodDefinition, "is_data_processing_function", staticmethod(lambda p: True) + ) + parameter_map = { + "params": Parameter( + "params", + Parameter.POSITIONAL_OR_KEYWORD, + annotation=Annotated[_MyQueryParams, "meta"], + ) + } + out = MethodDefinition.format_params("/x/y", parameter_map) + assert "data" in out + assert "start_date" in out + assert "tags" in out + + +def test_build_command_method_body_raw_query_params_fallback(monkeypatch): + from openbb_core.provider.abstract.query_params import QueryParams + + class _RawQueryParams(QueryParams): + symbol: str = "AAPL" + + monkeypatch.setattr( + MethodDefinition, "is_deprecated_function", staticmethod(lambda p: False) + ) + monkeypatch.setattr( + MethodDefinition, "query_params_model", staticmethod(lambda _a: None) + ) + + def endpoint(params: _RawQueryParams): + return params + + code = MethodDefinition.build_command_method_body("/x/y", endpoint) + assert '"symbol": symbol' in code diff --git a/openbb_platform/core/tests/app/static/test_container.py b/openbb_platform/core/tests/app/static/test_container.py index 6288b2544ec..4a97c3254e2 100644 --- a/openbb_platform/core/tests/app/static/test_container.py +++ b/openbb_platform/core/tests/app/static/test_container.py @@ -51,6 +51,25 @@ def test_container__run(mock_sync_run, container): mock_sync_run.assert_called_once() +@patch("openbb_core.app.command_runner.CommandRunner.sync_run") +def test_container__run_wraps_stream(mock_sync_run, container): + """Test container _run wraps a streaming response in an OBBStream.""" + from openbb_core.app.model.stream import OBBStream + + class _StreamingResponse: + media_type = "text/event-stream" + + def __init__(self): + self.body_iterator = self._gen() + + async def _gen(self): + yield "data: 0\n\n" + + mock_sync_run.return_value = _StreamingResponse() + result = container._run("/x/y") + assert isinstance(result, OBBStream) + + def test_container__check_credentials(container): """Test container _check_credentials method.""" assert container._check_credentials("provider_1") is False diff --git a/openbb_platform/core/tests/provider/abstract/test_fetcher_streaming.py b/openbb_platform/core/tests/provider/abstract/test_fetcher_streaming.py new file mode 100644 index 00000000000..135bf69a706 --- /dev/null +++ b/openbb_platform/core/tests/provider/abstract/test_fetcher_streaming.py @@ -0,0 +1,69 @@ +"""Tests for streaming fetchers.""" + +import asyncio +from collections.abc import AsyncIterator + +import pytest + +from openbb_core.provider.abstract.data import Data +from openbb_core.provider.abstract.fetcher import Fetcher +from openbb_core.provider.abstract.query_params import QueryParams + + +class _StreamQueryParams(QueryParams): + limit: int = 3 + + +class _StreamData(Data): + tick: int + + +class _StreamFetcher(Fetcher[_StreamQueryParams, AsyncIterator[_StreamData]]): + @staticmethod + def transform_query(params: dict) -> _StreamQueryParams: + return _StreamQueryParams(**params) + + @staticmethod + async def aextract_data(query, credentials, **kwargs) -> AsyncIterator[dict]: + async def _gen(): + for index in range(query.limit): + yield {"tick": index} + + return _gen() + + @staticmethod + async def stream_data(query, data, **kwargs) -> AsyncIterator[_StreamData]: + async for record in data: + yield _StreamData.model_validate(record) + + +def test_is_streaming(): + assert _StreamFetcher.is_streaming is True + + +def test_data_type_unwraps_async_iterator(): + assert _StreamFetcher.data_type is _StreamData + + +def test_fetch_data_returns_stream(): + async def _run(): + stream = await _StreamFetcher.fetch_data({"limit": 2}) + return [row async for row in stream] + + rows = asyncio.run(_run()) + assert [row.tick for row in rows] == [0, 1] + assert all(isinstance(row, _StreamData) for row in rows) + + +def test_fetcher_test_streaming_passes(): + _StreamFetcher.test({"limit": 3}) + + +def test_test_stream_handles_short_stream(): + query = _StreamFetcher.transform_query({"limit": 1}) + _StreamFetcher._test_stream(query, None, max_items=2) + + +def test_base_stream_data_raises(): + with pytest.raises(NotImplementedError): + Fetcher.stream_data(None, None) diff --git a/openbb_platform/core/tests/provider/standard_models/test_options_chains.py b/openbb_platform/core/tests/provider/standard_models/test_options_chains.py index e34d37c38ed..6f4988de773 100644 --- a/openbb_platform/core/tests/provider/standard_models/test_options_chains.py +++ b/openbb_platform/core/tests/provider/standard_models/test_options_chains.py @@ -280,6 +280,37 @@ class TestOpenAPISchema: ] assert any(r and r.endswith("OptionsChainsData") for r in refs) + def test_nested_model_field_emits_defs_in_openapi(self): + """A nested-model field on a subclass surfaces ``$defs`` in the schema. + + The row model's ``$defs`` (definitions for nested models referenced by + ``$ref``) are popped from the item schema and re-attached at the array + level so the OpenAPI document carries the nested definitions. The base + model has only scalar fields, so this path is exercised through a + subclass that adds a nested model. + """ + from fastapi import FastAPI + from pydantic import BaseModel, Field + + class _Greeks(BaseModel): + delta: float | None = None + + class _SubChains(OptionsChainsData): + greeks: list[_Greeks] | None = Field( + default=None, description="Greeks for the contract." + ) + + app = FastAPI() + + @app.get("/sub") + def sub() -> OBBject[_SubChains]: # pragma: no cover - schema-only + ... + + component = app.openapi()["components"]["schemas"]["_SubChains"] + assert component["type"] == "array" + assert "$defs" in component + assert "_Greeks" in component["$defs"] + def test_row_schema_field_set_pins_serializer_contract(self): """Snapshot guard: every model field appears in the row schema.