From db175b27520c4836d53af5873fd4c70da01e6c77 Mon Sep 17 00:00:00 2001 From: adityajunwal Date: Tue, 30 Dec 2025 19:29:27 +0530 Subject: [PATCH] api status --- src/db_services/ConfigurationServices.py | 59 ++++++++++++++++++- src/db_services/api_key_status_service.py | 19 ++++++ .../commonServices/baseService/baseService.py | 1 + src/services/commonServices/common.py | 16 +++-- src/services/utils/api_key_status_helper.py | 28 +++++++++ src/services/utils/common_utils.py | 14 ++--- src/services/utils/getConfiguration.py | 2 + 7 files changed, 123 insertions(+), 16 deletions(-) create mode 100644 src/db_services/api_key_status_service.py create mode 100644 src/services/utils/api_key_status_helper.py diff --git a/src/db_services/ConfigurationServices.py b/src/db_services/ConfigurationServices.py index 22f37ea2..509dd2cc 100644 --- a/src/db_services/ConfigurationServices.py +++ b/src/db_services/ConfigurationServices.py @@ -208,13 +208,13 @@ async def get_bridges_with_tools_and_apikeys(bridge_id, org_id, version_id=None) }, { '$project': { 'service': 1, 'apikey': 1, 'apikey_limit': { '$ifNull': ['$apikey_limit', 0] }, - 'apikey_usage': { '$ifNull': ['$apikey_usage', 0] } } + 'apikey_usage': { '$ifNull': ['$apikey_usage', 0] }, 'status': {'$ifNull': ['$status', None]} } } ], 'as': 'apikeys_docs' } }, - # Stage 5: Map each service to its corresponding apikey, handling empty case + # Stage 5: Map each service to its corresponding apikey and status, handling empty case { '$addFields': { 'apikeys': { @@ -267,9 +267,64 @@ async def get_bridges_with_tools_and_apikeys(bridge_id, org_id, version_id=None) }, {} ] + }, + 'apikey_status': { + '$cond': [ + { '$gt' : [ {'$size': '$apikeys_array'}, 0] }, + { + '$arrayToObject': { + '$map': { + 'input': '$apikeys_array', + 'as': 'item', + 'in': [ + '$$item.k', + { + '$let': { + 'vars': { + 'matched': { + '$arrayElemAt': [ + { + '$filter': { + 'input': '$apikeys_docs', + 'as': 'doc', + 'cond': { + '$eq': [ + '$$doc._id', + { + '$convert': { + 'input': '$$item.v', + 'to': 'objectId', + 'onError': None, + 'onNull': None + } + } + ] + } + } + }, + 0 + ] + } + }, + 'in': { + '$cond': [ + { '$ne' : ['$$matched', None] }, + '$$matched.status', + None + ] + } + } + } + ] + } + } + }, + {} + ] } } }, + # Stage 6: Lookup 'rag_parent_datas' using 'doc_ids' { '$lookup': { diff --git a/src/db_services/api_key_status_service.py b/src/db_services/api_key_status_service.py new file mode 100644 index 00000000..a7bb60de --- /dev/null +++ b/src/db_services/api_key_status_service.py @@ -0,0 +1,19 @@ +from bson import ObjectId +from globals import logger +from models.mongo_connection import db + +apikeyCredentialsModel = db["apikeycredentials"] + +async def update_apikey_status(apikey_id: str, status: str): + if not apikey_id: + return + + try: + result = await apikeyCredentialsModel.update_one( + {"_id": ObjectId(apikey_id)}, + {"$set": {"status": status}} + ) + if not result.modified_count: + logger.warning(f"No apikey credential updated for id={apikey_id}") + except Exception as exc: + logger.error(f"Failed to update API key status for {apikey_id}: {exc}") \ No newline at end of file diff --git a/src/services/commonServices/baseService/baseService.py b/src/services/commonServices/baseService/baseService.py index 9860b232..28fb94fe 100644 --- a/src/services/commonServices/baseService/baseService.py +++ b/src/services/commonServices/baseService/baseService.py @@ -56,6 +56,7 @@ def __init__(self, params): self.type = params.get('type') self.token_calculator = params.get('token_calculator') self.apikey_object_id = params.get('apikey_object_id') + self.apikey_status = params.get('apikey_status') self.image_data = params.get('images') self.tool_call_count = params.get('tool_call_count') self.text = params.get('text') diff --git a/src/services/commonServices/common.py b/src/services/commonServices/common.py index d161c67d..415fb3cc 100644 --- a/src/services/commonServices/common.py +++ b/src/services/commonServices/common.py @@ -38,8 +38,8 @@ process_background_tasks_for_playground, process_variable_state, handle_agent_transfer, - update_cost_and_last_used_in_background, - setup_agent_pre_tools + update_cost_usage_and_apikey_status_in_background, + setup_agent_pre_tools, ) from src.services.utils.guardrails_validator import guardrails_check from src.services.utils.rich_text_support import process_chatbot_response @@ -122,6 +122,8 @@ async def chat_multiple_agents(request_body): async def chat(request_body): result ={} class_obj= {} + first_execution_error_code = None + completion_success = True try: # Store bridge_configurations for potential transfer logic bridge_configurations = request_body.get('body', {}).get('bridge_configurations', {}) @@ -243,6 +245,7 @@ async def chat(request_body): # Handle exceptions during execution execution_failed = True original_error = str(execution_exception) + first_execution_error_code = original_error.split()[7] original_exception = execution_exception logger.error(f"Initial execution failed with {parsed_data['service']}/{parsed_data['model']}: {original_error}") result = { @@ -367,8 +370,8 @@ async def chat(request_body): result['response']['testcase_result'] = testcase_result else: await process_background_tasks_for_playground(result, parsed_data) - await update_cost_and_last_used_in_background(parsed_data) + # Save agent bridge_id to Redis for 3 days (259200 seconds) thread_id = parsed_data.get('thread_id') sub_thread_id = parsed_data.get('sub_thread_id') @@ -420,10 +423,11 @@ async def chat(request_body): } if parsed_data['is_playground'] and parsed_data['body']['bridge_configurations'].get('playground_response_format'): await sendResponse(parsed_data['body']['bridge_configurations']['playground_response_format'], error_object, success=False, variables=parsed_data.get('variables',{})) + + completion_success = False raise ValueError(error_object) - - - + finally: + await update_cost_usage_and_apikey_status_in_background(parsed_data, first_execution_error_code, completion_success) @handle_exceptions async def orchestrator_chat(request_body): try: diff --git a/src/services/utils/api_key_status_helper.py b/src/services/utils/api_key_status_helper.py new file mode 100644 index 00000000..7c8194df --- /dev/null +++ b/src/services/utils/api_key_status_helper.py @@ -0,0 +1,28 @@ +from src.db_services.api_key_status_service import update_apikey_status + +STATUS_BY_CODE = { + "401": "invalid", + "429": "exhuasted", +} + +def classify_status_from_error(code) -> str: + if code in STATUS_BY_CODE: + return STATUS_BY_CODE[code] + + return code + +async def mark_apikey_status_from_response(parsed_data, code=None): + apikey_map = parsed_data.get("apikey_object_id") or {} + status_map = parsed_data.get("apikey_status") or {} + service = parsed_data.get("service") + apikey_id = apikey_map.get(service) + if not apikey_id: + return + + new_status = "working" if not code else classify_status_from_error( + code + ) + if status_map.get(service) == new_status: + return # already up to date; skip DB write + + await update_apikey_status(apikey_id, new_status) \ No newline at end of file diff --git a/src/services/utils/common_utils.py b/src/services/utils/common_utils.py index 29aaa75c..90d21fb8 100644 --- a/src/services/utils/common_utils.py +++ b/src/services/utils/common_utils.py @@ -30,6 +30,7 @@ from ..commonServices.baseService.utils import sendResponse from src.services.utils.rich_text_support import process_chatbot_response from src.db_services.orchestrator_history_service import orchestrator_collector +from src.services.utils.api_key_status_helper import mark_apikey_status_from_response def setup_agent_pre_tools(parsed_data, bridge_configurations): """ @@ -164,6 +165,7 @@ def parse_request_body(request_body): "usage" : {}, "type" : body.get('configuration',{}).get('type'), "apikey_object_id" : body.get('apikey_object_id'), + "apikey_status": body.get('apikey_status'), "images" : body.get('images'), "tool_call_count": body.get('tool_call_count'), "tokens" : {}, @@ -1294,12 +1296,8 @@ async def update_cost_and_last_used(parsed_data): except Exception as e: logger.error(f"Error updating cost and last used: {str(e)}") -async def update_cost_and_last_used_in_background(parsed_data): - """Kick off the async cost cache update using the data available on parsed_data.""" - if not isinstance(parsed_data, dict): - logger.warning("Skipping background cost update due to invalid parsed data.") - return - - asyncio.create_task(update_cost_and_last_used(parsed_data)) - +async def update_cost_usage_and_apikey_status_in_background(parsed_data, code, completion_success): + if completion_success: + asyncio.create_task(update_cost_and_last_used(parsed_data)) + asyncio.create_task(mark_apikey_status_from_response(parsed_data, code)) \ No newline at end of file diff --git a/src/services/utils/getConfiguration.py b/src/services/utils/getConfiguration.py index a80250b9..50bd6265 100644 --- a/src/services/utils/getConfiguration.py +++ b/src/services/utils/getConfiguration.py @@ -75,6 +75,7 @@ async def _prepare_configuration_response(configuration, service, bridge_id, api apikey = setup_api_key(service, result, apikey, chatbot) apikey_object_id = result.get('bridges', {}).get('apikey_object_id') + apikey_status = result.get('bridges', {}).get('apikey_status') # Handle image type early return if configuration['type'] == 'image': @@ -148,6 +149,7 @@ async def _prepare_configuration_response(configuration, service, bridge_id, api 'service': service, 'apikey': apikey, 'apikey_object_id': apikey_object_id, + 'apikey_status': apikey_status, 'RTLayer': RTLayer, 'template': template_content.get('template') if template_content else None, 'user_reference': result.get('bridges', {}).get('user_reference', ''),