# API 密钥管理方案

本文档概述了 `airs_api_keys` 表的数据库存储、查询和代码建议方案，以支持 LLM 和搜索引擎 API 密钥的灵活管理。

## 1. `airs_api_keys` 表的原始结构

```sql
create table public.airs_api_keys (
  id bigint generated by default as identity not null,
  created_at timestamp with time zone not null default now(),
  api_key text null,
  ran_at timestamp without time zone null default now(),
  provider_id bigint null,
  constraint airs_api_keys_pkey primary key (id),
  constraint airs_api_keys_provider_id_fkey foreign KEY (provider_id) references airs_api_providers (id)
) TABLESPACE pg_default;
```

## 2. `airs_api_keys` 表的建议修改结构

为了支持不同类型的 API 密钥（LLM 和搜索引擎）以及更灵活的查找机制，建议修改 `airs_api_keys` 表结构如下：

```sql
create table public.airs_api_keys (
  id bigint generated by default as identity not null,
  created_at timestamp with time zone not null default now(),
  api_key text null,
  ran_at timestamp without time zone null default now(),
  provider_id bigint null,
  api_type text not null, -- 'llm' or 'search'
  lookup_key text not null, -- LLM 的模型名称 (e.g., 'glm-4', 'gemini-pro'), 搜索引擎的域名或特定路径 (e.g., 'api.tavily.com')
  url_pattern text null, -- 可选: 用于验证或更精确匹配的 URL 模式 (e.g., 'https://open.bigmodel.cn/api/paas/v4/chat/completions')
  constraint airs_api_keys_pkey primary key (id),
  constraint airs_api_keys_provider_id_fkey foreign KEY (provider_id) references airs_api_providers (id),
  constraint airs_api_keys_unique_lookup_key unique (api_type, lookup_key) -- 确保每种类型和查找键组合的唯一性
) TABLESPACE pg_default;
```

**字段说明：**

* `api_type`: 字符串类型，用于区分 API 密钥的类型，例如 `'llm'` 或 `'search'`。
* `lookup_key`: 字符串类型，作为查询 API 密钥的关键字段。
  * 对于 LLM 类 API，这将是请求体中 `model` 字段的值（例如 `glm-4`, `gemini-pro`）。
  * 对于搜索引擎类 API，这将是其域名或特定路径（例如 `api.tavily.com`）。
* `url_pattern`: 可选的字符串类型，用于存储与 API 密钥关联的完整或部分 URL 模式。这可以用于在某些情况下进行额外的验证或更精确的匹配。
* `constraint airs_api_keys_unique_lookup_key`: 确保 `api_type` 和 `lookup_key` 的组合是唯一的，防止重复的 API 密钥配置。

### 示例数据库记录

以下是一些 `airs_api_keys` 表中可能存在的示例记录：

| id  | created_at             | api_key                    | ran_at              | provider_id | api_type | lookup_key     | url_pattern                                                              |
| --- | ---------------------- | -------------------------- | ------------------- | ----------- | -------- | -------------- | ------------------------------------------------------------------------ |
| 1   | 2023-01-01 10:00:00+00 | sk-xxxxxxxxxxxxxxxxxxxxglm | 2023-01-01 10:00:00 | 1           | llm      | glm-4          | https://open.bigmodel.cn/api/paas/v4/chat/completions                    |
| 2   | 2023-01-02 11:00:00+00 | AIzaSyBxxxxxxxxxxxxxxxxxx  | 2023-01-02 11:00:00 | 2           | llm      | gemini-pro     | https://generativelanguage.googleapis.com/v1beta/openai/chat/completions |
| 3   | 2023-01-03 12:00:00+00 | tvly-xxxxxxxxxxxxxxxxxxxx  | 2023-01-03 12:00:00 | 3           | search   | api.tavily.com | https://api.tavily.com/search                                            |

## 3. 数据库查询方案

根据修改后的表结构，可以按 `api_type` 和 `lookup_key` 高效地查询 API 密钥。

### 查询 LLM API 密钥

当代理服务收到 LLM 请求时，它会从请求体中提取 `model` 名称，并使用它作为 `lookup_key` 进行查询。

```sql
SELECT api_key, provider_id
FROM public.airs_api_keys
WHERE api_type = 'llm' AND lookup_key = :model_name;
```

* `:model_name` 将替换为从请求体中解析出的模型名称，例如 `'glm-4'` 或 `'gemini-pro'`。

### 查询搜索引擎 API 密钥

当代理服务收到搜索引擎请求时，它会根据请求的 URL 识别出对应的域名或特定路径，并使用它作为 `lookup_key` 进行查询。

```sql
SELECT api_key, provider_id
FROM public.airs_api_keys
WHERE api_type = 'search' AND lookup_key = :domain_or_path;
```

* `:domain_or_path` 将替换为从请求 URL 中提取的域名或路径，例如 `'api.tavily.com'`。

## 4. 代码建议方案

为了在 FastAPI 代理服务中实现 API 密钥的动态获取和注入，需要修改 `api_key_sb.py` 和 `proxy.py` 文件。

### `api_key_sb.py` 中的 `get_api_key_info` 函数

修改 `get_api_key_info` 函数以接受 `api_type` 和 `lookup_value` 作为参数，并根据这些参数从 Supabase 查询 API 密钥。

```python
# api_key_sb.py

from supabase import create_client, Client
import os
from dotenv import load_dotenv

load_dotenv()

SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

async def get_api_key_info(api_type: str, lookup_value: str):
    """
    根据 API 类型和查找值从 Supabase 获取 API 密钥信息。
    """
    try:
        response = supabase.from_('airs_api_keys').select('api_key, provider_id').eq('api_type', api_type).eq('lookup_key', lookup_value).limit(1).execute()
        if response.data:
            return response.data[0]
        return None
    except Exception as e:
        print(f"Error fetching API key info: {e}")
        return None
```

### `proxy.py` 中的请求处理逻辑

在 `proxy.py` 中，需要修改请求处理逻辑以：

1. 识别传入请求的 API 类型（LLM 或搜索引擎）。
2. 根据 API 类型从请求中提取相应的 `lookup_value`（`model` 名称或域名/路径）。
3. 调用 `get_api_key_info` 获取 API 密钥。
4. 将获取到的 API 密钥注入到转发请求的相应位置（例如，请求头或请求体）。

以下是 `proxy.py` 中 `handle_proxy_request` 函数的伪代码片段：

```python
# proxy.py

from fastapi import FastAPI, Request, Response, HTTPException
from fastapi.responses import StreamingResponse
import httpx
import json
import logging
import re
from api_key_sb import get_api_key_info # 导入新的函数

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()
client = httpx.AsyncClient()

# 统一错误响应格式
def create_error_response(status_code: int, message: str):
    return Response(
        content=json.dumps({"error": {"message": message, "type": "proxy_error"}}),
        status_code=status_code,
        media_type="application/json"
    )

@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH", "TRACE"])
async def handle_proxy_request(request: Request, path: str):
    target_url_base = path # 假设 path 已经包含了目标 URL 的基础部分
    api_key = None
    api_type = None
    lookup_value = None

    # 识别 API 类型并提取 lookup_value
    if "open.bigmodel.cn" in target_url_base or "generativelanguage.googleapis.com" in target_url_base:
        api_type = "llm"
        try:
            request_body = await request.json()
            lookup_value = request_body.get("model")
            if not lookup_value:
                logger.warning(f"LLM request missing 'model' in body for path: {path}")
                return create_error_response(400, "LLM request body must contain 'model' field.")
        except json.JSONDecodeError:
            logger.warning(f"LLM request body is not valid JSON for path: {path}")
            return create_error_response(400, "LLM request body must be valid JSON.")
    elif "api.tavily.com" in target_url_base:
        api_type = "search"
        lookup_value = "api.tavily.com" # 或者从 URL 中更动态地提取，例如使用正则表达式
    else:
        logger.warning(f"Unknown API type for path: {path}")
        return create_error_response(400, "Unknown API type or unsupported target URL.")

    # 从数据库获取 API 密钥
    if api_type and lookup_value:
        api_key_info = await get_api_key_info(api_type, lookup_value)
        if api_key_info:
            api_key = api_key_info["api_key"]
            logger.info(f"Successfully retrieved API key for type: {api_type}, lookup_value: {lookup_value}")
        else:
            logger.warning(f"API Key not found for type: {api_type}, lookup_value: {lookup_value}")
            return create_error_response(401, "API Key not found for the requested service.")
    else:
        logger.error(f"Failed to determine API type or lookup value for path: {path}")
        return create_error_response(500, "Internal server error: Could not determine API key parameters.")

    # 构建目标 URL
    # 假设原始请求的 path 已经包含了完整的后端服务 URL，例如 /v1/https/open.bigmodel.cn/api/paas/v4/chat/completions
    # 我们需要提取出 https://open.bigmodel.cn/api/paas/v4/chat/completions
    match = re.match(r"/v1/(https?://.*)", path)
    if not match:
        logger.error(f"Invalid target URL format in path: {path}")
        return create_error_response(400, "Invalid target URL format.")

    full_target_url = match.group(1)

    # 准备转发请求的头部和内容
    headers = dict(request.headers)
    headers.pop("host", None) # 移除 host 头，避免后端服务解析错误
    headers.pop("authorization", None) # 移除原始的 authorization 头，我们将注入新的 API 密钥

    # 注入 API 密钥
    if api_type == "llm":
        # 对于 LLM，通常 API 密钥在 Authorization 头中
        headers["Authorization"] = f"Bearer {api_key}"
    elif api_type == "search":
        # 对于 Tavily 等搜索引擎，API 密钥可能在请求体中或自定义头中
        # 这里假设 Tavily API 密钥在请求体中，需要修改请求体
        # 注意：修改请求体需要重新构建请求，这可能比较复杂
        # 更简单的做法是如果后端支持，通过自定义头传递
        # 假设 Tavily API 密钥在请求体中，且请求体是 JSON
        if request.method == "POST":
            try:
                request_body = await request.json()
                request_body["api_key"] = api_key # 注入 Tavily API 密钥
                request_content = json.dumps(request_body).encode("utf-8")
                headers["Content-Type"] = "application/json"
            except json.JSONDecodeError:
                logger.error(f"Search request body is not valid JSON for path: {path}")
                return create_error_response(400, "Search request body must be valid JSON.")
        else:
            # 如果是 GET 请求，API 密钥可能在查询参数中，这里需要根据实际情况调整
            logger.warning(f"Search API key injection for GET request not implemented for path: {path}")
            return create_error_response(501, "Search API key injection for GET requests is not yet supported.")

    # 转发请求
    try:
        req = client.build_request(
            method=request.method,
            url=full_target_url,
            headers=headers,
            content=request_content if 'request_content' in locals() else await request.body(),
            params=request.query_params,
            timeout=30.0 # 设置超时时间
        )

        proxy_response = await client.send(req, stream=True)

        return StreamingResponse(
            proxy_response.aiter_bytes(),
            status_code=proxy_response.status_code,
            headers=proxy_response.headers
        )
    except httpx.RequestError as e:
        logger.error(f"HTTPX Request Error for {full_target_url}: {e}")
        return create_error_response(500, f"Proxy request failed: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred: {e}")
        return create_error_response(500, "An unexpected error occurred during proxying.")
```