| 1 | from typing import Dict |
| 2 | from typing import List |
| 3 | from typing import Union |
| 4 | |
| 5 | from cortex4py.api import Api |
| 6 | from fastapi import HTTPException |
| 7 | from loguru import logger |
| 8 | |
| 9 | from app.connectors.cortex.schema.analyzers import AnalyzerJobData |
| 10 | from app.connectors.cortex.schema.analyzers import AnalyzersResponse |
| 11 | from app.connectors.cortex.schema.analyzers import RunAnalyzerBody |
| 12 | from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse |
| 13 | from app.connectors.cortex.utils.universal import ( |
| 14 | create_cortex_client, # Importing create_cortex_client; Importing from universal.py |
| 15 | ) |
| 16 | from app.connectors.cortex.utils.universal import run_and_wait_for_analyzer |
| 17 | |
| 18 | |
| 19 | async def fetch_analyzers(api: Api) -> List[Dict]: |
| 20 | """ |
| 21 | Fetches all analyzers from the API. |
| 22 | |
| 23 | Args: |
| 24 | api (Api): The API object used to make the request. |
| 25 | |
| 26 | Returns: |
| 27 | List[Dict]: A list of dictionaries representing the analyzers. |
| 28 | """ |
| 29 | try: |
| 30 | return api.analyzers.find_all({}, range="all") |
| 31 | except Exception as e: |
| 32 | logger.error(f"Error fetching analyzers: {e}") |
| 33 | raise HTTPException(status_code=500, detail=f"Error fetching analyzers: {e}") |
| 34 | |
| 35 | |
| 36 | def extract_analyzer_names(analyzers: List[Dict]) -> List[str]: |
| 37 | """ |
| 38 | Extracts the names of analyzers from a list of dictionaries. |
| 39 | |
| 40 | Args: |
| 41 | analyzers (List[Dict]): A list of dictionaries representing analyzers. |
| 42 | |
| 43 | Returns: |
| 44 | List[str]: A list of analyzer names. |
| 45 | |
| 46 | Raises: |
| 47 | HTTPException: If there is an error processing the analyzers. |
| 48 | """ |
| 49 | try: |
| 50 | return [analyzer.name for analyzer in analyzers] |
| 51 | except Exception as e: |
| 52 | logger.error(f"Error processing analyzers: {e}") |
| 53 | raise HTTPException(status_code=500, detail=f"Error processing analyzers: {e}") |
| 54 | |
| 55 | |
| 56 | async def init_cortex_client() -> Union[Api, None]: |
| 57 | """ |
| 58 | Initializes the Cortex client. |
| 59 | |
| 60 | Returns: |
| 61 | Union[Api, None]: The initialized Cortex client or None if initialization fails. |
| 62 | """ |
| 63 | return await create_cortex_client("Cortex") |
| 64 | |
| 65 | |
| 66 | def handle_api_initialization(api: Union[Api, None]) -> Api: |
| 67 | """ |
| 68 | Handles the initialization of the API. |
| 69 | |
| 70 | Args: |
| 71 | api (Union[Api, None]): The API object to be initialized. |
| 72 | |
| 73 | Returns: |
| 74 | Api: The initialized API object. |
| 75 | |
| 76 | Raises: |
| 77 | HTTPException: If the API initialization fails. |
| 78 | """ |
| 79 | if api is None: |
| 80 | logger.error("API initialization failed") |
| 81 | raise HTTPException(status_code=500, detail="API initialization failed") |
| 82 | return api |
| 83 | |
| 84 | |
| 85 | async def get_analyzers() -> AnalyzersResponse: |
| 86 | """ |
| 87 | Fetches the list of analyzers from the Cortex API. |
| 88 | |
| 89 | Returns: |
| 90 | AnalyzersResponse: The response object containing the list of analyzer names. |
| 91 | """ |
| 92 | api = await init_cortex_client() |
| 93 | handle_api_initialization(api) |
| 94 | |
| 95 | analyzers = await fetch_analyzers(api) |
| 96 | analyzer_names = extract_analyzer_names(analyzers) |
| 97 | |
| 98 | return AnalyzersResponse( |
| 99 | success=True, |
| 100 | message="Successfully fetched analyzers", |
| 101 | analyzers=analyzer_names, |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | async def run_analyzer( |
| 106 | run_analyzer_body: RunAnalyzerBody, |
| 107 | data_type: str, |
| 108 | ) -> RunAnalyzerResponse: |
| 109 | """ |
| 110 | Runs an analyzer with the given analyzer name, analyzer data, and data type. |
| 111 | |
| 112 | Args: |
| 113 | run_analyzer_body (RunAnalyzerBody): The body of the run analyzer request. |
| 114 | data_type (str): The type of data being analyzed. |
| 115 | |
| 116 | Returns: |
| 117 | RunAnalyzerResponse: The response containing the success status, message, and analyzer report. |
| 118 | """ |
| 119 | api = await init_cortex_client() |
| 120 | handle_api_initialization(api) |
| 121 | |
| 122 | analyzer_name = run_analyzer_body.analyzer_name |
| 123 | analyzer_data = run_analyzer_body.analyzer_data |
| 124 | logger.info( |
| 125 | f"Running analyzer {analyzer_name} with data {analyzer_data} of type {data_type}", |
| 126 | ) |
| 127 | job_data = AnalyzerJobData(data=analyzer_data, dataType=data_type) |
| 128 | |
| 129 | result = await run_and_wait_for_analyzer( |
| 130 | analyzer_name=analyzer_name, |
| 131 | job_data=job_data, |
| 132 | ) |
| 133 | |
| 134 | if result is None: |
| 135 | logger.error(f"Failed to run analyzer {analyzer_name}") |
| 136 | raise HTTPException( |
| 137 | status_code=500, |
| 138 | detail=f"Failed to run analyzer {analyzer_name}", |
| 139 | ) |
| 140 | |
| 141 | return RunAnalyzerResponse( |
| 142 | success=True, |
| 143 | message="Successfully ran analyzer", |
| 144 | report=result, |
| 145 | ) |