Skip to content

Commit

Permalink
feat:Improve API Key Error Handling (#428)
Browse files Browse the repository at this point in the history
* 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")
  • Loading branch information
XiangZhang-zx authored Jan 20, 2025
1 parent 3b0036f commit 69f0e48
Showing 1 changed file with 16 additions and 6 deletions.
22 changes: 16 additions & 6 deletions aios/llm_core/adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down

0 comments on commit 69f0e48

Please sign in to comment.