Fixes for async filter search endpoint

pyproject.toml

@@ -129,7 +129,7 @@ server = [
     "packaging==22.0",
     "pyparsing==3.2.1",
     "setuptools<81",
-    "supervisor==4.1.0",
+    "supervisor==4.3.*",
     "uwsgi==2.0.23",
 ]
 ansible = ["ansible==14.4.*"]

usaspending_api/common/helpers/decorators.py

@@ -1,14 +1,23 @@
 import logging
+from typing import Callable, Coroutine
 
+from asgiref.sync import sync_to_async
 from django.db import connection
 from django.db.utils import OperationalError
+from django.http import HttpResponseBase, StreamingHttpResponse
+from ninja import Router
+from ninja.decorators import decorate_view
+from rest_framework.request import Request
+from rest_framework.response import Response as DRFResponse
+from rest_framework.views import APIView
 
 from usaspending_api.common.exceptions import EndpointTimeoutException
+from usaspending_api.common.helpers.endpoint_documentation import DJANGO_NINJA_TEMP_NAME_FOR_DRF
 
 logger = logging.getLogger(__name__)
 
 
-def set_db_timeout(timeout_in_seconds):
+def set_db_timeout(timeout_in_seconds: int) -> Callable:
     """Decorator used to set the database statement timeout within the Django app scope
 
     Args:
@@ -32,8 +41,8 @@ def func_running_db_call(...):
     """
     timeout_in_ms = int(timeout_in_seconds * 1000)
 
-    def wrap(func):
-        def wrapper(*args, **kwargs):
+    def wrap(func: Callable) -> Callable:
+        def wrapper(*args, **kwargs) -> Callable:
             with connection.cursor() as cursor:
                 cursor.execute("show statement_timeout")
                 prev_timeout = cursor.fetchall()[0][0]
@@ -57,8 +66,10 @@ def wrapper(*args, **kwargs):
 
             try:
                 func_response = func(*args, **kwargs)
-            except OperationalError:
-                raise EndpointTimeoutException("Django ORM exceeded the specified timeout of %ds" % timeout_in_seconds)
+            except OperationalError as exc:
+                raise EndpointTimeoutException(
+                    "Django ORM exceeded the specified timeout of %ds" % timeout_in_seconds
+                ) from exc
             finally:
                 with connection.cursor() as cursor:
                     cursor.execute("show statement_timeout")
@@ -85,3 +96,67 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return wrap
+
+
+def _show_browsable_result(self: APIView, request: Request, *args, **kwargs) -> DRFResponse:
+    """Renders the real Ninja response that the browsable()-installed run() wrapper
+    already buffered onto the request, through DRF's normal browsable rendering pipeline."""
+    try:
+        status_code, body_text = request._browsable_result
+    except AttributeError as exc:
+        raise NotImplementedError("browsable() did not buffer a result for this request.") from exc
+    return DRFResponse(body_text, status=status_code)
+
+
+async def _buffer_streaming_response(response: HttpResponseBase) -> tuple[int, str]:
+    if isinstance(response, StreamingHttpResponse):
+        body = b"".join([part async for part in response])
+    else:
+        body = response.content
+    return response.status_code, body.decode()
+
+
+def browsable(
+    router: Router, path: str, *, endpoint_doc: str, methods: tuple[str] = ("POST",), **ninja_kwargs
+) -> Callable:
+    """
+    This decorator is a temporary solution while we have the DRF API UI that needs to display the endpoints
+    defined with Django Ninja.
+    """
+
+    def decorator(view_func: Callable) -> Callable:
+        doc_view = type(
+            f"{view_func.__name__.title().replace('_', '')}",
+            (APIView,),
+            {
+                "__doc__": view_func.__doc__,
+                "endpoint_doc": endpoint_doc,
+                **{m.lower(): _show_browsable_result for m in methods},
+            },
+        )
+        rendered = doc_view.as_view()
+        router.get(path, auth=None, include_in_schema=False, url_name=DJANGO_NINJA_TEMP_NAME_FOR_DRF)(
+            lambda request: rendered(request)
+        )
+
+        registered_view_func = router.api_operation(list(methods), path, **ninja_kwargs)(view_func)
+        operation = registered_view_func._ninja_operation
+        if not operation.is_async:
+            raise RuntimeError("@browsable only supports async Ninja operations.")
+
+        def _html_aware_run(run_func: Callable) -> Callable:
+            async def wrapper(request: Request, *args, **kwargs) -> Coroutine:
+                result = await run_func(request, *args, **kwargs)
+                if "text/html" not in request.headers.get("Accept", ""):
+                    return result
+                status_code, body_text = await _buffer_streaming_response(result)
+                request._browsable_result = (status_code, body_text)
+                return await sync_to_async(rendered, thread_sensitive=True)(request)
+
+            return wrapper
+
+        decorate_view(_html_aware_run)(registered_view_func)
+
+        return registered_view_func
+
+    return decorator

usaspending_api/common/helpers/endpoint_documentation.py

@@ -1,5 +1,9 @@
+import inspect
 import os
 import re
+from typing import Callable
+
+from django.urls import URLPattern
 
 from usaspending_api.settings import REPO_DIR
 
@@ -20,8 +24,12 @@
     "/api/v2/bulk_download/list_database_download_files",
 ]
 
+# A constant value used to distinguish the views created to satisfy display Django Ninja views
+# in the current Django REST Framework API UI
+DJANGO_NINJA_TEMP_NAME_FOR_DRF = "for_django_ninja_in_drf"
+
 
-def case_sensitive_file_exists(file_path):
+def case_sensitive_file_exists(file_path: str) -> bool:
     """
     File names are case insensitive on Macs and Windows and case sensitive on
     Linux.  This is one way to perform case sensitive file checking on a case
@@ -31,7 +39,7 @@ def case_sensitive_file_exists(file_path):
     return os.path.isfile(file_path) and filename in os.listdir(directory)
 
 
-def get_endpoint_urls_doc_paths_and_docstrings(endpoint_prefixes=None):
+def get_endpoint_urls_doc_paths_and_docstrings(endpoint_prefixes: str | None = None) -> list[tuple[str, URLPattern]]:
     """
     Compiles the list of endpoint URLs and associated RegexURLPattern objects
     serviced by Django.  If path_prefixes is supplied, returned URLS will be
@@ -44,19 +52,23 @@ def get_endpoint_urls_doc_paths_and_docstrings(endpoint_prefixes=None):
 
     results = []
 
-    def _traverse_urls(base, url_patterns):
+    def _traverse_urls(base: str, url_patterns: list[URLPattern], resolver_name: str | None = None) -> None:
         for url in url_patterns:
+            if resolver_name == "ninja" and getattr(url, "name", None) != DJANGO_NINJA_TEMP_NAME_FOR_DRF:
+                # For the purpose of the DRF docs on Django Ninja endpoints we only want the views that are created
+                # by the "browsable" decorator (see usaspending_api.common.helpers.decorators.browsable)
+                continue
             cleaned = base + url.pattern.regex.pattern.lstrip("^").rstrip("$")
             if hasattr(url, "url_patterns"):
-                _traverse_urls(cleaned, url.url_patterns)
+                _traverse_urls(cleaned, url.url_patterns, url.app_name)
             elif not endpoint_prefixes or cleaned.startswith(endpoint_prefixes):
                 results.append((cleaned, url))
 
     _traverse_urls("/", urls.urlpatterns)
     return results
 
 
-def get_endpoints_from_endpoints_markdown():
+def get_endpoints_from_endpoints_markdown() -> list[str]:
     """
     Looks for and extracts URLs for patterns like |[display](url)|method|description|
     from the master endpoints.md markdown file.
@@ -66,7 +78,7 @@ def get_endpoints_from_endpoints_markdown():
     return [e.split("?")[0] for e in ENDPOINT_PATTERN.findall(contents) if e]
 
 
-def get_fully_qualified_name(obj):
+def get_fully_qualified_name(obj: object) -> str:
     """
     Fully qualifies an object name.  For example, would return
     "usaspending_api.common.helpers.endpoint_documentation.get_fully_qualified_name"
@@ -75,7 +87,7 @@ def get_fully_qualified_name(obj):
     return "{}.{}".format(obj.__module__, obj.__qualname__)
 
 
-def validate_docs(url, url_object, master_endpoint_list):
+def validate_docs(url: str, url_object: URLPattern, master_endpoint_list: list[str]) -> list[str]:  # noqa: PLR0912
     """
     Ensures that an endpoint_doc property and a docstring is provided for the
     view associated with the provided url and checks that the URL is mentioned
@@ -88,15 +100,15 @@ def validate_docs(url, url_object, master_endpoint_list):
         return []
 
     # Handles class and function based views.
-    view = url_object.callback.cls if getattr(url_object.callback, "cls", None) else url_object.callback
+    view = get_view_for_endpoint(url, url_object)
 
     messages = []
 
     if not hasattr(view, "endpoint_doc"):
         if url not in _EXCLUDED_URLS:
             messages.append("{} ({}) missing endpoint_doc property".format(qualified_name, url))
     else:
-        endpoint_doc = getattr(view, "endpoint_doc")
+        endpoint_doc = view.endpoint_doc
         if not endpoint_doc:
             messages.append("{}.endpoint_doc ({}) is invalid".format(qualified_name, url))
         else:
@@ -127,3 +139,40 @@ def validate_docs(url, url_object, master_endpoint_list):
             messages.append("No URL found in {} that matches {} ({})".format(ENDPOINTS_MD, url, qualified_name))
 
     return messages
+
+
+def get_view_for_endpoint(url: str, url_object: URLPattern) -> type:
+    view = url_object.callback.cls if getattr(url_object.callback, "cls", None) else url_object.callback
+
+    # Handle the possibility that "endpoint_doc" doesn't exist in the expected location because it was
+    # implemented with django-ninja
+    if not hasattr(view, "endpoint_doc") and url not in _EXCLUDED_URLS:
+        view = unwrap_ninja_view(view)
+
+    return view
+
+
+def unwrap_ninja_view(callback: type | Callable) -> type:
+    """
+    In order for Django Ninja endpoints to appear in the DRF API UI a new view is created that renders onto the DRF UI.
+    To handle this new view in a similar way as the current DRF views we need to step through the decorators and
+    retrieve the view class that has the docstring and "endpoint_doc" property added dynamically.
+    """
+    try:
+        path_view = inspect.getclosurevars(callback).nonlocals.get("self")
+    except TypeError:
+        return callback
+
+    operations = getattr(path_view, "operations", None)
+    if not operations:
+        return callback
+
+    view_func = operations[0].view_func
+
+    if not hasattr(view_func, "cls"):
+        try:
+            closure_vars = inspect.getclosurevars(view_func).nonlocals
+        except TypeError:
+            closure_vars = {}
+        view_func = next((v for v in closure_vars.values() if hasattr(v, "cls")), view_func)
+    return getattr(view_func, "cls", view_func)

usaspending_api/common/management/commands/check_for_endpoint_documentation.py

@@ -1,17 +1,23 @@
+import logging
+
 from django.core.management.base import BaseCommand
+
 from usaspending_api.common.helpers.endpoint_documentation import (
+    _EXCLUDED_URLS,
     CURRENT_ENDPOINT_PREFIXES,
     get_endpoint_urls_doc_paths_and_docstrings,
     get_endpoints_from_endpoints_markdown,
+    get_view_for_endpoint,
     validate_docs,
 )
 
+logger = logging.getLogger(__name__)
 
-class Command(BaseCommand):
 
+class Command(BaseCommand):
     help = "Checks endpoints to ensure they point to documentation and are contained in the master endpoint list."
 
-    def handle(self, *args, **options):
+    def handle(self, *args, **options) -> None:
         master_endpoint_list = get_endpoints_from_endpoints_markdown()
 
         endpoints = get_endpoint_urls_doc_paths_and_docstrings(CURRENT_ENDPOINT_PREFIXES)
@@ -22,20 +28,20 @@ def handle(self, *args, **options):
 
         if messages:
             for message in messages:
-                print(message)
+                logger.error(message)
             exit(1)
 
         # By request, let's count how much documentation falls in the contracts directory vs not-contracts.
         contract_count = sum(
-            getattr(
-                e[1].callback.cls if getattr(e[1].callback, "cls", None) else e[1].callback, "endpoint_doc", ""
-            ).startswith("usaspending_api/api_contracts/contracts/")
+            getattr(get_view_for_endpoint(*e), "endpoint_doc", "").startswith(
+                "usaspending_api/api_contracts/contracts/"
+            )
             for e in endpoints
         )
 
-        print("Looks like endpoint documentation is happy, healthy, wealthy, and wise.  Current tally:")
-        print("    contract count: {:,}".format(contract_count))
-        print("    non-contract count: {:,}".format(len(endpoints) - contract_count))
+        logger.info("Looks like endpoint documentation is happy, healthy, wealthy, and wise.  Current tally:")
+        logger.info(f"    contract count: {contract_count:,}")
+        logger.info(f"    non-contract count: {len(endpoints) - (contract_count + len(_EXCLUDED_URLS)):,}")
+        logger.info(f"    excluded count: {len(_EXCLUDED_URLS):,}")
         if (len(endpoints) - contract_count) == 0:
-            print()
-            print("*** GOOD JOB TEAM! ***")
+            logger.info("\n*** GOOD JOB TEAM! ***")

usaspending_api/llm/assistants/filter_search.py

@@ -9,7 +9,7 @@
 
 from usaspending_api.common.helpers.aws_helpers import async_aws_client
 from usaspending_api.llm.models.db_models import Assistant, Message, Session, ToolUse
-from usaspending_api.llm.models.py_models import AITool, build_defc_description
+from usaspending_api.llm.models.py_models import AITool, abuild_defc_description
 
 logger = logging.getLogger(__name__)
 
@@ -29,6 +29,7 @@ def __init__(self, assistant: Assistant, tools: list[AITool], session: Session)
         self.message_order = 0
         self.messages = []
         self.tool_iterations = 0
+        self._tool_config: dict[str, list[dict]] | None = None
 
     @staticmethod
     def _extract_text_from_content(content: list[dict]) -> str:
@@ -111,14 +112,17 @@ async def _create_message_from_response(self, response: dict) -> Message:
         self.messages.append(output_message)
         return message
 
-    @cached_property
-    def tool_config(self) -> dict[str, list[dict]]:
-        specs = [tool.description.model_dump() for tool in self.tools]
-        for spec in specs:
-            def_codes_property = spec["input_schema"].get("properties", {}).get("defCodes")
-            if def_codes_property is not None:
-                def_codes_property["description"] = build_defc_description()
-        return {"tools": [{"toolSpec": {"inputSchema": {"json": spec.pop("input_schema")}, **spec}} for spec in specs]}
+    async def aget_tool_config(self) -> dict[str, list[dict]]:
+        if self._tool_config is None:
+            specs = [tool.description.model_dump() for tool in self.tools]
+            for spec in specs:
+                def_codes_property = spec["input_schema"].get("properties", {}).get("defCodes")
+                if def_codes_property is not None:
+                    def_codes_property["description"] = await abuild_defc_description()
+            self._tool_config = {
+                "tools": [{"toolSpec": {"inputSchema": {"json": spec.pop("input_schema")}, **spec}} for spec in specs]
+            }
+        return self._tool_config
 
     @staticmethod
     def _fiscal_year_date_context() -> str:
@@ -182,7 +186,7 @@ async def search(self, query: str) -> AsyncGenerator[dict[str, str], None]:
             response = await client.converse(
                 modelId=self.assistant.ai_model.model_id,
                 messages=self.messages,
-                toolConfig=self.tool_config,
+                toolConfig=await self.aget_tool_config(),
                 system=[{"text": self.system_message}],
                 inferenceConfig=self.inference_config,
             )
@@ -222,7 +226,7 @@ async def search(self, query: str) -> AsyncGenerator[dict[str, str], None]:
                 response = await client.converse(
                     modelId=self.assistant.ai_model.model_id,
                     messages=self.messages,
-                    toolConfig=self.tool_config,
+                    toolConfig=await self.aget_tool_config(),
                     system=[{"text": self.system_message}],
                     inferenceConfig=self.inference_config,
                 )