mirror of
https://github.com/BlueSkyXN/AI2API.git
synced 2026-04-22 15:47:19 +08:00
0.2.0
This commit is contained in:
BlueSkyXN
@@ -0,0 +1,131 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
import urllib.parse
|
||||
from typing import Dict, Any
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def base64url_encode(data: bytes) -> str:
|
||||
"""将数据进行 base64url 编码"""
|
||||
return base64.urlsafe_b64encode(data).decode('utf-8').replace('=', '')
|
||||
|
||||
|
||||
def sha256_hash(input_data: bytes) -> bytes:
|
||||
"""计算输入数据的 SHA-256 哈希值"""
|
||||
return hashlib.sha256(input_data).digest()
|
||||
|
||||
|
||||
def create_oauth_state() -> Dict[str, Any]:
|
||||
"""创建 OAuth 状态对象"""
|
||||
code_verifier_bytes = secrets.token_bytes(32)
|
||||
code_verifier = base64url_encode(code_verifier_bytes)
|
||||
|
||||
code_challenge_bytes = sha256_hash(code_verifier.encode('utf-8'))
|
||||
code_challenge = base64url_encode(code_challenge_bytes)
|
||||
|
||||
state = base64url_encode(secrets.token_bytes(8))
|
||||
|
||||
oauth_state = {
|
||||
"codeVerifier": code_verifier,
|
||||
"codeChallenge": code_challenge,
|
||||
"state": state,
|
||||
"creationTime": int(import_time())
|
||||
}
|
||||
|
||||
return oauth_state
|
||||
|
||||
|
||||
def generate_authorize_url(oauth_state: Dict[str, Any]) -> str:
|
||||
"""生成授权 URL"""
|
||||
client_id = "v"
|
||||
|
||||
params = {
|
||||
"response_type": "code",
|
||||
"code_challenge": oauth_state["codeChallenge"],
|
||||
"client_id": client_id,
|
||||
"state": oauth_state["state"],
|
||||
"prompt": "login"
|
||||
}
|
||||
|
||||
query_string = urllib.parse.urlencode(params)
|
||||
authorize_url = f"https://auth.augmentcode.com/authorize?{query_string}"
|
||||
|
||||
return authorize_url
|
||||
|
||||
|
||||
def get_access_token(tenant_url: str, code_verifier: str, code: str) -> str:
|
||||
"""获取访问令牌"""
|
||||
data = {
|
||||
"grant_type": "authorization_code",
|
||||
"client_id": "v",
|
||||
"code_verifier": code_verifier,
|
||||
"redirect_uri": "",
|
||||
"code": code
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
f"{tenant_url}token",
|
||||
json=data,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
json_response = response.json()
|
||||
token = json_response.get("access_token")
|
||||
|
||||
return token
|
||||
|
||||
|
||||
def parse_code(code_str: str) -> Dict[str, str]:
|
||||
"""解析返回的代码"""
|
||||
parsed = json.loads(code_str)
|
||||
return {
|
||||
"code": parsed.get("code"),
|
||||
"state": parsed.get("state"),
|
||||
"tenant_url": parsed.get("tenant_url")
|
||||
}
|
||||
|
||||
|
||||
def import_time():
|
||||
"""获取当前时间戳"""
|
||||
import time
|
||||
return time.time() * 1000
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("正在生成 OAuth 状态...")
|
||||
oauth_state = create_oauth_state()
|
||||
print("OAuth 状态已生成:")
|
||||
print(json.dumps(oauth_state, indent=2))
|
||||
|
||||
url = generate_authorize_url(oauth_state)
|
||||
print("\n请访问以下 URL 进行授权:")
|
||||
print(url)
|
||||
|
||||
code_str = input("\n请输入返回的代码 (JSON 格式): ")
|
||||
|
||||
try:
|
||||
parsed_code = parse_code(code_str)
|
||||
print("代码已解析:")
|
||||
print(json.dumps(parsed_code, indent=2))
|
||||
|
||||
print("\n正在获取访问令牌...")
|
||||
token = get_access_token(
|
||||
parsed_code["tenant_url"],
|
||||
oauth_state["codeVerifier"],
|
||||
parsed_code["code"]
|
||||
)
|
||||
|
||||
print("\n访问令牌:")
|
||||
print(token)
|
||||
|
||||
except Exception as e:
|
||||
print(f"发生错误: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,744 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
OpenAI to Augment API Adapter
|
||||
|
||||
这个FastAPI应用程序将OpenAI API请求格式转换为Augment API格式，
|
||||
允许OpenAI客户端直接与Augment服务通信。
|
||||
所有配置参数都通过命令行参数提供，不依赖于环境变量或配置文件。
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import uuid
|
||||
import time
|
||||
import logging
|
||||
import argparse
|
||||
from typing import List, Optional, Dict, Any, Literal, Union
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
from fastapi import FastAPI, Header, HTTPException, Depends, Request
|
||||
from fastapi.responses import StreamingResponse, JSONResponse
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
import uvicorn
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#################################################
|
||||
# 模型定义
|
||||
#################################################
|
||||
|
||||
# OpenAI API 请求模型
|
||||
class ChatMessage(BaseModel):
|
||||
"""表示OpenAI聊天API中的单条消息"""
|
||||
role: Literal["system", "user", "assistant", "function"]
|
||||
content: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
|
||||
class ChatCompletionRequest(BaseModel):
|
||||
"""OpenAI聊天完成API请求模型"""
|
||||
model: str
|
||||
messages: List[ChatMessage]
|
||||
temperature: Optional[float] = 1.0
|
||||
top_p: Optional[float] = 1.0
|
||||
n: Optional[int] = 1
|
||||
stream: Optional[bool] = False
|
||||
max_tokens: Optional[int] = None
|
||||
presence_penalty: Optional[float] = 0
|
||||
frequency_penalty: Optional[float] = 0
|
||||
user: Optional[str] = None
|
||||
|
||||
# OpenAI API 响应模型
|
||||
class ChatCompletionResponseChoice(BaseModel):
|
||||
"""OpenAI聊天完成API响应中的单个选择"""
|
||||
index: int
|
||||
message: ChatMessage
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
class Usage(BaseModel):
|
||||
"""OpenAI API响应中的token使用信息"""
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
total_tokens: int
|
||||
|
||||
class ChatCompletionResponse(BaseModel):
|
||||
"""OpenAI聊天完成API响应模型"""
|
||||
id: str
|
||||
object: str = "chat.completion"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[ChatCompletionResponseChoice]
|
||||
usage: Usage
|
||||
|
||||
# OpenAI API 流式响应模型
|
||||
class ChatCompletionStreamResponseChoice(BaseModel):
|
||||
"""OpenAI聊天完成流式API响应中的单个选择"""
|
||||
index: int
|
||||
delta: Dict[str, Any]
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
class ChatCompletionStreamResponse(BaseModel):
|
||||
"""OpenAI聊天完成流式API响应模型"""
|
||||
id: str
|
||||
object: str = "chat.completion.chunk"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[ChatCompletionStreamResponseChoice]
|
||||
|
||||
# 模型信息响应
|
||||
class ModelInfo(BaseModel):
|
||||
"""OpenAI模型信息"""
|
||||
id: str
|
||||
object: str = "model"
|
||||
created: int
|
||||
owned_by: str = "augment"
|
||||
|
||||
class ModelListResponse(BaseModel):
|
||||
"""OpenAI模型列表响应"""
|
||||
object: str = "list"
|
||||
data: List[ModelInfo]
|
||||
|
||||
# Augment API 请求相关模型
|
||||
class AugmentResponseNode(BaseModel):
|
||||
"""Augment API响应节点"""
|
||||
id: int
|
||||
type: int
|
||||
content: str
|
||||
tool_use: Optional[Any] = None
|
||||
|
||||
class AugmentChatHistoryItem(BaseModel):
|
||||
"""Augment API聊天历史记录条目"""
|
||||
request_message: str
|
||||
response_text: str
|
||||
request_id: Optional[str] = None
|
||||
request_nodes: List[Any] = []
|
||||
response_nodes: List[AugmentResponseNode] = []
|
||||
|
||||
class AugmentBlobs(BaseModel):
|
||||
"""Augment API Blobs对象"""
|
||||
checkpoint_id: Optional[str] = None
|
||||
added_blobs: List[Any] = []
|
||||
deleted_blobs: List[Any] = []
|
||||
|
||||
class AugmentVcsChange(BaseModel):
|
||||
"""Augment API VCS更改"""
|
||||
working_directory_changes: List[Any] = []
|
||||
|
||||
class AugmentFeatureFlags(BaseModel):
|
||||
"""Augment API功能标志"""
|
||||
support_raw_output: bool = True
|
||||
|
||||
# 完整的Augment API请求模型
|
||||
class AugmentChatRequest(BaseModel):
|
||||
"""Augment API聊天请求模型 - 基于抓包分析更新"""
|
||||
model: Optional[str] = None
|
||||
path: Optional[str] = None
|
||||
prefix: Optional[str] = None
|
||||
selected_code: Optional[str] = None
|
||||
suffix: Optional[str] = None
|
||||
message: str
|
||||
chat_history: List[AugmentChatHistoryItem] = []
|
||||
lang: Optional[str] = None
|
||||
blobs: AugmentBlobs = AugmentBlobs()
|
||||
user_guided_blobs: List[Any] = []
|
||||
context_code_exchange_request_id: Optional[str] = None
|
||||
vcs_change: AugmentVcsChange = AugmentVcsChange()
|
||||
recency_info_recent_changes: List[Any] = []
|
||||
external_source_ids: List[Any] = []
|
||||
disable_auto_external_sources: Optional[bool] = None
|
||||
user_guidelines: str = ""
|
||||
workspace_guidelines: str = ""
|
||||
feature_detection_flags: AugmentFeatureFlags = AugmentFeatureFlags()
|
||||
tool_definitions: List[Any] = []
|
||||
nodes: List[Any] = []
|
||||
mode: str = "CHAT"
|
||||
agent_memories: Optional[Any] = None
|
||||
system_prompt: Optional[str] = None # 保留此字段以兼容之前的代码
|
||||
|
||||
# Augment API响应模型
|
||||
class AugmentResponseChunk(BaseModel):
|
||||
"""Augment API响应块"""
|
||||
text: str
|
||||
unknown_blob_names: List[Any] = []
|
||||
checkpoint_not_found: bool = False
|
||||
workspace_file_chunks: List[Any] = []
|
||||
incorporated_external_sources: List[Any] = []
|
||||
nodes: List[AugmentResponseNode] = []
|
||||
|
||||
#################################################
|
||||
# 辅助函数
|
||||
#################################################
|
||||
|
||||
def generate_id():
|
||||
"""生成唯一ID，类似于OpenAI的格式"""
|
||||
return str(uuid.uuid4()).replace("-", "")[:24]
|
||||
|
||||
def estimate_tokens(text):
|
||||
"""
|
||||
估计文本的token数量
|
||||
这是一个简单的估算，实际数量可能有所不同
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
# 简单估算：假设每个单词约等于1.3个token
|
||||
# 中文字符每个字约等于1个token
|
||||
words = len(text.split()) if text else 0
|
||||
chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff') if text else 0
|
||||
return int(words * 1.3 + chinese_chars)
|
||||
|
||||
def convert_to_augment_request(openai_request: ChatCompletionRequest) -> AugmentChatRequest:
|
||||
"""
|
||||
将OpenAI API请求转换为Augment API请求
|
||||
|
||||
Args:
|
||||
openai_request: OpenAI API请求对象
|
||||
|
||||
Returns:
|
||||
转换后的Augment API请求对象
|
||||
|
||||
Raises:
|
||||
HTTPException: 如果请求格式无效
|
||||
"""
|
||||
chat_history = []
|
||||
system_message = None
|
||||
|
||||
# 处理消息历史记录
|
||||
for i in range(len(openai_request.messages) - 1):
|
||||
msg = openai_request.messages[i]
|
||||
if msg.role == "system":
|
||||
system_message = msg.content
|
||||
elif msg.role == "user" and i + 1 < len(openai_request.messages) and openai_request.messages[i + 1].role == "assistant":
|
||||
user_msg = msg.content
|
||||
assistant_msg = openai_request.messages[i + 1].content
|
||||
|
||||
# 创建历史记录条目，格式符合Augment API
|
||||
history_item = AugmentChatHistoryItem(
|
||||
request_message=user_msg,
|
||||
response_text=assistant_msg,
|
||||
request_id=generate_id(),
|
||||
response_nodes=[
|
||||
AugmentResponseNode(
|
||||
id=0,
|
||||
type=0,
|
||||
content=assistant_msg,
|
||||
tool_use=None
|
||||
)
|
||||
]
|
||||
)
|
||||
chat_history.append(history_item)
|
||||
|
||||
# 获取当前用户消息
|
||||
current_message = None
|
||||
for msg in reversed(openai_request.messages):
|
||||
if msg.role == "user":
|
||||
current_message = msg.content
|
||||
break
|
||||
|
||||
# 如果没有用户消息，则返回错误
|
||||
if current_message is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="At least one user message is required"
|
||||
)
|
||||
|
||||
# 准备Augment请求体
|
||||
augment_request = AugmentChatRequest(
|
||||
message=current_message,
|
||||
chat_history=chat_history,
|
||||
mode="CHAT"
|
||||
)
|
||||
|
||||
# 如果有系统消息，设置为用户指南
|
||||
if system_message:
|
||||
augment_request.user_guidelines = system_message
|
||||
|
||||
return augment_request
|
||||
|
||||
#################################################
|
||||
# FastAPI应用
|
||||
#################################################
|
||||
|
||||
def create_app(augment_base_url, chat_endpoint, timeout):
|
||||
"""
|
||||
创建并配置FastAPI应用
|
||||
|
||||
Args:
|
||||
augment_base_url: Augment API基础URL
|
||||
chat_endpoint: 聊天端点路径
|
||||
timeout: 请求超时时间
|
||||
|
||||
Returns:
|
||||
配置好的FastAPI应用
|
||||
"""
|
||||
app = FastAPI(
|
||||
title="OpenAI to Augment API Adapter",
|
||||
description="A FastAPI adapter that converts OpenAI API requests to Augment API format",
|
||||
version="1.0.0"
|
||||
)
|
||||
|
||||
# 添加CORS中间件
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
#################################################
|
||||
# 中间件和依赖项
|
||||
#################################################
|
||||
|
||||
@app.middleware("http")
|
||||
async def catch_exceptions_middleware(request: Request, call_next):
|
||||
"""捕获所有未处理的异常，返回适当的错误响应"""
|
||||
try:
|
||||
return await call_next(request)
|
||||
except Exception as e:
|
||||
logger.exception("Unhandled exception")
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={
|
||||
"error": {
|
||||
"message": str(e),
|
||||
"type": "internal_server_error",
|
||||
"param": None,
|
||||
"code": "internal_server_error"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
async def verify_api_key(authorization: str = Header(...)):
|
||||
"""
|
||||
验证API密钥
|
||||
|
||||
Args:
|
||||
authorization: Authorization头部值
|
||||
|
||||
Returns:
|
||||
提取的API密钥
|
||||
|
||||
Raises:
|
||||
HTTPException: 如果API密钥格式无效或为空
|
||||
"""
|
||||
if not authorization.startswith("Bearer "):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Invalid API key format. Expected 'Bearer YOUR_API_KEY'",
|
||||
"type": "invalid_request_error",
|
||||
"param": "authorization",
|
||||
"code": "invalid_api_key"
|
||||
}
|
||||
}
|
||||
)
|
||||
api_key = authorization.replace("Bearer ", "")
|
||||
if not api_key:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "API key cannot be empty",
|
||||
"type": "invalid_request_error",
|
||||
"param": "authorization",
|
||||
"code": "invalid_api_key"
|
||||
}
|
||||
}
|
||||
)
|
||||
return api_key
|
||||
|
||||
#################################################
|
||||
# API端点
|
||||
#################################################
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
"""健康检查端点"""
|
||||
return {"status": "ok", "timestamp": datetime.now().isoformat()}
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def list_models():
|
||||
"""列出支持的模型"""
|
||||
# 返回一个虚拟的模型列表
|
||||
models = [
|
||||
ModelInfo(id="gpt-3.5-turbo", created=int(time.time())),
|
||||
ModelInfo(id="gpt-4", created=int(time.time())),
|
||||
ModelInfo(id="augment-default", created=int(time.time())),
|
||||
]
|
||||
return ModelListResponse(data=models)
|
||||
|
||||
@app.get("/v1/models/{model_id}")
|
||||
async def get_model(model_id: str):
|
||||
"""获取特定模型的信息"""
|
||||
return ModelInfo(id=model_id, created=int(time.time()))
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat_completions(
|
||||
request: ChatCompletionRequest,
|
||||
api_key: str = Depends(verify_api_key)
|
||||
):
|
||||
"""
|
||||
聊天完成端点 - 将OpenAI API请求转换为Augment API请求
|
||||
|
||||
Args:
|
||||
request: OpenAI格式的聊天完成请求
|
||||
api_key: 通过验证的API密钥
|
||||
|
||||
Returns:
|
||||
OpenAI格式的聊天完成响应或流式响应
|
||||
"""
|
||||
try:
|
||||
# 转换为Augment请求格式
|
||||
augment_request = convert_to_augment_request(request)
|
||||
logger.debug(f"Converted request: {augment_request.dict()}")
|
||||
|
||||
# 决定是否使用流式响应
|
||||
if request.stream:
|
||||
return StreamingResponse(
|
||||
stream_augment_response(augment_base_url, api_key, augment_request, request.model, chat_endpoint, timeout),
|
||||
media_type="text/event-stream"
|
||||
)
|
||||
else:
|
||||
# 同步请求处理
|
||||
return await handle_sync_request(augment_base_url, api_key, augment_request, request.model, chat_endpoint, timeout)
|
||||
|
||||
except httpx.TimeoutException:
|
||||
logger.error("Request to Augment API timed out")
|
||||
raise HTTPException(
|
||||
status_code=504,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Request to Augment API timed out",
|
||||
"type": "timeout_error",
|
||||
"param": None,
|
||||
"code": "timeout"
|
||||
}
|
||||
}
|
||||
)
|
||||
except httpx.HTTPError as e:
|
||||
logger.error(f"HTTP error: {str(e)}")
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Error communicating with Augment API: {str(e)}",
|
||||
"type": "api_error",
|
||||
"param": None,
|
||||
"code": "api_error"
|
||||
}
|
||||
}
|
||||
)
|
||||
except HTTPException:
|
||||
# 重新抛出HTTPException，以保持原始状态码和详细信息
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.exception("Unexpected error")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Internal server error: {str(e)}",
|
||||
"type": "internal_server_error",
|
||||
"param": None,
|
||||
"code": "internal_server_error"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
async def handle_sync_request(base_url, api_key, augment_request, model_name, chat_endpoint, timeout):
|
||||
"""
|
||||
处理同步请求
|
||||
|
||||
Args:
|
||||
base_url: Augment API基础URL
|
||||
api_key: API密钥
|
||||
augment_request: Augment API请求对象
|
||||
model_name: 模型名称
|
||||
chat_endpoint: 聊天端点
|
||||
timeout: 请求超时时间
|
||||
|
||||
Returns:
|
||||
OpenAI格式的聊天完成响应
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
response = await client.post(
|
||||
f"{base_url.rstrip('/')}/{chat_endpoint}",
|
||||
json=augment_request.dict(),
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"User-Agent": "Augment.openai-adapter/1.0.0",
|
||||
"Accept": "*/*"
|
||||
}
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.error(f"Augment API error: {response.status_code} - {response.text}")
|
||||
raise HTTPException(
|
||||
status_code=response.status_code,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Augment API error: {response.text}",
|
||||
"type": "api_error",
|
||||
"param": None,
|
||||
"code": "api_error"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# 处理流式响应，合并为完整响应
|
||||
full_response = ""
|
||||
for line in response.text.split("\n"):
|
||||
if line.strip():
|
||||
try:
|
||||
data = json.loads(line)
|
||||
if "text" in data and data["text"]:
|
||||
full_response += data["text"]
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse JSON: {line}")
|
||||
|
||||
# 估算token使用情况
|
||||
prompt_tokens = estimate_tokens(augment_request.message)
|
||||
completion_tokens = estimate_tokens(full_response)
|
||||
|
||||
# 构建OpenAI格式响应
|
||||
return ChatCompletionResponse(
|
||||
id=f"chatcmpl-{generate_id()}",
|
||||
created=int(time.time()),
|
||||
model=model_name,
|
||||
choices=[
|
||||
ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=ChatMessage(
|
||||
role="assistant",
|
||||
content=full_response
|
||||
),
|
||||
finish_reason="stop"
|
||||
)
|
||||
],
|
||||
usage=Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens
|
||||
)
|
||||
)
|
||||
|
||||
async def stream_augment_response(base_url, api_key, augment_request, model_name, chat_endpoint, timeout):
|
||||
"""
|
||||
处理流式响应
|
||||
|
||||
Args:
|
||||
base_url: Augment API基础URL
|
||||
api_key: API密钥
|
||||
augment_request: Augment API请求对象
|
||||
model_name: 模型名称
|
||||
chat_endpoint: 聊天端点
|
||||
timeout: 请求超时时间
|
||||
|
||||
Yields:
|
||||
流式响应的数据块
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
try:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
f"{base_url.rstrip('/')}/{chat_endpoint}",
|
||||
json=augment_request.dict(),
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"User-Agent": "Augment.openai-adapter/1.0.0",
|
||||
"Accept": "*/*"
|
||||
}
|
||||
) as response:
|
||||
|
||||
if response.status_code != 200:
|
||||
error_detail = await response.aread()
|
||||
logger.error(f"Augment API error: {response.status_code} - {error_detail}")
|
||||
error_message = f"Error from Augment API: {error_detail.decode('utf-8', errors='replace')}"
|
||||
yield f"data: {json.dumps({'error': error_message})}\n\n"
|
||||
return
|
||||
|
||||
# 生成唯一ID
|
||||
chat_id = f"chatcmpl-{generate_id()}"
|
||||
created_time = int(time.time())
|
||||
|
||||
# 初始化响应
|
||||
init_response = ChatCompletionStreamResponse(
|
||||
id=chat_id,
|
||||
created=created_time,
|
||||
model=model_name,
|
||||
choices=[
|
||||
ChatCompletionStreamResponseChoice(
|
||||
index=0,
|
||||
delta={"role": "assistant"},
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
init_data = json.dumps(init_response.dict())
|
||||
yield f"data: {init_data}\n\n"
|
||||
|
||||
# 处理流式响应
|
||||
buffer = ""
|
||||
async for line in response.aiter_lines():
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
# 解析Augment响应格式
|
||||
chunk = json.loads(line)
|
||||
if "text" in chunk and chunk["text"]:
|
||||
content = chunk["text"]
|
||||
|
||||
# 发送增量更新
|
||||
stream_response = ChatCompletionStreamResponse(
|
||||
id=chat_id,
|
||||
created=created_time,
|
||||
model=model_name,
|
||||
choices=[
|
||||
ChatCompletionStreamResponseChoice(
|
||||
index=0,
|
||||
delta={"content": content},
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
response_data = json.dumps(stream_response.dict())
|
||||
yield f"data: {response_data}\n\n"
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse JSON: {line}")
|
||||
|
||||
# 发送完成信号
|
||||
final_response = ChatCompletionStreamResponse(
|
||||
id=chat_id,
|
||||
created=created_time,
|
||||
model=model_name,
|
||||
choices=[
|
||||
ChatCompletionStreamResponseChoice(
|
||||
index=0,
|
||||
delta={},
|
||||
finish_reason="stop"
|
||||
)
|
||||
]
|
||||
)
|
||||
final_data = json.dumps(final_response.dict())
|
||||
yield f"data: {final_data}\n\n"
|
||||
|
||||
# 发送[DONE]标记
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
except httpx.TimeoutException:
|
||||
logger.error("Request to Augment API timed out")
|
||||
yield f"data: {json.dumps({'error': 'Request to Augment API timed out'})}\n\n"
|
||||
except httpx.HTTPError as e:
|
||||
logger.error(f"HTTP error: {str(e)}")
|
||||
yield f"data: {json.dumps({'error': f'Error communicating with Augment API: {str(e)}'})}\n\n"
|
||||
except Exception as e:
|
||||
logger.exception("Unexpected error")
|
||||
yield f"data: {json.dumps({'error': f'Internal server error: {str(e)}'})}\n\n"
|
||||
|
||||
def parse_args():
|
||||
"""解析命令行参数"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="OpenAI to Augment API Adapter",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--augment-url",
|
||||
default="https://d18.api.augmentcode.com/",
|
||||
help="Augment API基础URL"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--chat-endpoint",
|
||||
default="chat-stream",
|
||||
help="Augment聊天端点路径"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
default="0.0.0.0",
|
||||
help="服务器主机地址"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8686,
|
||||
help="服务器端口"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
type=int,
|
||||
default=120,
|
||||
help="API请求超时时间（秒）"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
action="store_true",
|
||||
help="启用调试模式"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--tenant-id",
|
||||
default="d18",
|
||||
help="Augment API租户ID (域名前缀)"
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
#################################################
|
||||
# 主程序
|
||||
#################################################
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
args = parse_args()
|
||||
|
||||
# 配置日志级别
|
||||
if args.debug:
|
||||
logging.getLogger().setLevel(logging.DEBUG)
|
||||
|
||||
# 构建完整的Augment URL
|
||||
if args.augment_url == "https://d18.api.augmentcode.com/":
|
||||
# 如果使用默认URL，则应用tenant-id参数
|
||||
augment_base_url = f"https://{args.tenant_id}.api.augmentcode.com/"
|
||||
logger.info(f"Using tenant ID: {args.tenant_id}")
|
||||
else:
|
||||
# 否则使用提供的URL
|
||||
augment_base_url = args.augment_url
|
||||
|
||||
# 创建应用
|
||||
app = create_app(
|
||||
augment_base_url=augment_base_url,
|
||||
chat_endpoint=args.chat_endpoint,
|
||||
timeout=args.timeout
|
||||
)
|
||||
|
||||
# 启动应用
|
||||
logger.info(f"Starting server on {args.host}:{args.port}")
|
||||
logger.info(f"Using Augment base URL: {augment_base_url}")
|
||||
logger.info(f"Using Augment chat endpoint: {args.chat_endpoint}")
|
||||
|
||||
uvicorn.run(
|
||||
app,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level="info" if not args.debug else "debug"
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user