From 69f0e489bb0673b263539e23b005aa2a8591c5dd Mon Sep 17 00:00:00 2001
From: xiang <129304623+XiangZhang-zx@users.noreply.github.com>
Date: Mon, 20 Jan 2025 13:28:07 -0500
Subject: [PATCH] feat:Improve API Key Error Handling (#428)

* Fix

* Feat: Prioritize API Key from Config File

Problem: The system was using environment variables for API keys.
Solution: Changed the code to first check the config file for API keys. If not found, it falls back to environment variables.
Impact: Ensures API keys are managed centrally in the config file, improving flexibility and control.

* Update adapter.py

* fix: use PR code instead of main branch in workflow

Before: Workflow always used main branch code
After: Workflow uses PR code when testing PR

* Update test_ollama.yml

* Update test_ollama.yml

* Update test_ollama.yml

* Update README: Add AIOS Refresh Command

* feat: Add dynamic LLM configuration support in AIOS kernel

# Description
This PR enables AIOS kernel to dynamically handle LLM configurations from users, allowing:

1. Dynamic API key updates for different LLM providers
2. Runtime configuration refresh without restart
3. Secure API key management

## Changes
- Added config update endpoint in kernel
- Added configuration refresh mechanism
- Updated environment variable handling

* Update kernel.py

* Fix:LLM API Key Error Handling

## What's Changed
- Added proper error handling for invalid API keys
- Fixed response parsing for LiteLLM completion calls
- Added HTTP status codes for different error types:
  - 402 for API key issues
  - 500 for other errors

* feat:Improve API Key Error Handling

- Added API key masking in error messages (e.g., "ab****12")
---
 aios/llm_core/adapter.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/aios/llm_core/adapter.py b/aios/llm_core/adapter.py
index feec7028..0f6420c9 100644
--- a/aios/llm_core/adapter.py
+++ b/aios/llm_core/adapter.py
@@ -306,18 +306,28 @@ def address_syscall(
                             temperature=temperature,
                         )
                 except Exception as e:
-                    # Handle API key related errors
-                    if "Invalid API key" in str(e) or "API key not found" in str(e):
+                    error_msg = str(e)
+                    # Mask API key in error message - only show first and last 2 chars
+                    if "API key provided:" in error_msg:
+                        key_start = error_msg.find("API key provided:") + len("API key provided: ")
+                        key_end = error_msg.find(".", key_start)
+                        if key_end == -1:  # If no period found, find next space
+                            key_end = error_msg.find(" ", key_start)
+                        if key_end != -1:  # If we found the end of the key
+                            api_key = error_msg[key_start:key_end]
+                            masked_key = f"{api_key[:2]}****{api_key[-2:]}" if len(api_key) > 4 else "****"
+                            error_msg = error_msg[:key_start] + masked_key + error_msg[key_end:]
+
+                    if "Invalid API key" in error_msg or "API key not found" in error_msg:
                         return Response(
                             response_message="Error: Invalid or missing API key for the selected model.",
-                            error=str(e),
+                            error=error_msg,
                             finished=True,
                             status_code=402
                         )
-                    # Handle other types of errors 
                     return Response(
-                        response_message=f"LLM Error: {str(e)}",
-                        error=str(e),
+                        response_message=f"LLM Error: {error_msg}",
+                        error=error_msg,
                         finished=True,
                         status_code=500
                     )