Update app.py

app.py

@@ -30,6 +30,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+
 #################################################
 # 模型定义
 #################################################
@@ -41,6 +42,7 @@ class ChatMessage(BaseModel):
     content: Optional[str] = None
     name: Optional[str] = None
 
+
 class ChatCompletionRequest(BaseModel):
     """OpenAI聊天完成API请求模型"""
     model: str
@@ -54,6 +56,7 @@ class ChatCompletionRequest(BaseModel):
     frequency_penalty: Optional[float] = 0
     user: Optional[str] = None
 
+
 # OpenAI API 响应模型
 class ChatCompletionResponseChoice(BaseModel):
     """OpenAI聊天完成API响应中的单个选择"""
@@ -61,12 +64,14 @@ class ChatCompletionResponseChoice(BaseModel):
     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
@@ -76,6 +81,7 @@ class ChatCompletionResponse(BaseModel):
     choices: List[ChatCompletionResponseChoice]
     usage: Usage
 
+
 # OpenAI API 流式响应模型
 class ChatCompletionStreamResponseChoice(BaseModel):
     """OpenAI聊天完成流式API响应中的单个选择"""
@@ -83,6 +89,7 @@ class ChatCompletionStreamResponseChoice(BaseModel):
     delta: Dict[str, Any]
     finish_reason: Optional[str] = None
 
+
 class ChatCompletionStreamResponse(BaseModel):
     """OpenAI聊天完成流式API响应模型"""
     id: str
@@ -91,6 +98,7 @@ class ChatCompletionStreamResponse(BaseModel):
     model: str
     choices: List[ChatCompletionStreamResponseChoice]
 
+
 # 模型信息响应
 class ModelInfo(BaseModel):
     """OpenAI模型信息"""
@@ -99,11 +107,13 @@ class ModelInfo(BaseModel):
     created: int
     owned_by: str = "augment"
 
+
 class ModelListResponse(BaseModel):
     """OpenAI模型列表响应"""
     object: str = "list"
     data: List[ModelInfo]
 
+
 # Augment API 请求相关模型
 class AugmentResponseNode(BaseModel):
     """Augment API响应节点"""
@@ -112,6 +122,7 @@ class AugmentResponseNode(BaseModel):
     content: str
     tool_use: Optional[Any] = None
 
+
 class AugmentChatHistoryItem(BaseModel):
     """Augment API聊天历史记录条目"""
     request_message: str
@@ -120,20 +131,24 @@ class AugmentChatHistoryItem(BaseModel):
     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聊天请求模型 - 基于抓包分析更新"""
@@ -157,10 +172,11 @@ class AugmentChatRequest(BaseModel):
     feature_detection_flags: AugmentFeatureFlags = AugmentFeatureFlags()
     tool_definitions: List[Any] = []
     nodes: List[Any] = []
-    mode: str = "CHAT"
+    mode: str = "AGENT"
     agent_memories: Optional[Any] = None
     system_prompt: Optional[str] = None  # 保留此字段以兼容之前的代码
 
+
 # Augment API响应模型
 class AugmentResponseChunk(BaseModel):
     """Augment API响应块"""
@@ -171,6 +187,7 @@ class AugmentResponseChunk(BaseModel):
     incorporated_external_sources: List[Any] = []
     nodes: List[AugmentResponseNode] = []
 
+
 #################################################
 # 辅助函数
 #################################################
@@ -179,6 +196,7 @@ def generate_id():
     """生成唯一ID，类似于OpenAI的格式"""
     return str(uuid.uuid4()).replace("-", "")[:24]
 
+
 def estimate_tokens(text):
     """
     估计文本的token数量
@@ -192,31 +210,33 @@ def estimate_tokens(text):
     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":
+        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,
@@ -232,34 +252,35 @@ def convert_to_augment_request(openai_request: ChatCompletionRequest) -> Augment
                 ]
             )
             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应用
 #################################################
@@ -267,12 +288,12 @@ def convert_to_augment_request(openai_request: ChatCompletionRequest) -> Augment
 def create_app(augment_base_url, chat_endpoint, timeout):
     """
     创建并配置FastAPI应用
-    
+
     Args:
         augment_base_url: Augment API基础URL
         chat_endpoint: 聊天端点路径
         timeout: 请求超时时间
-        
+
     Returns:
         配置好的FastAPI应用
     """
@@ -317,19 +338,19 @@ def create_app(augment_base_url, chat_endpoint, timeout):
     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, 
+                status_code=401,
                 detail={
                     "error": {
                         "message": "Invalid API key format. Expected 'Bearer YOUR_API_KEY'",
@@ -381,16 +402,16 @@ def create_app(augment_base_url, chat_endpoint, timeout):
 
     @app.post("/v1/chat/completions")
     async def chat_completions(
-        request: ChatCompletionRequest,
-        api_key: str = Depends(verify_api_key)
+            request: ChatCompletionRequest,
+            api_key: str = Depends(verify_api_key)
     ):
         """
         聊天完成端点 - 将OpenAI API请求转换为Augment API请求
-        
+
         Args:
             request: OpenAI格式的聊天完成请求
             api_key: 通过验证的API密钥
-            
+
         Returns:
             OpenAI格式的聊天完成响应或流式响应
         """
@@ -399,7 +420,6 @@ def create_app(augment_base_url, chat_endpoint, timeout):
             augment_request = convert_to_augment_request(request)
             logger.debug(f"Converted request: {augment_request.dict()}")
 
-
             if ":" in api_key:
                 tenant_id, api_key = api_key.split(":")
                 augment_base_url = f"https://{tenant_id}.api.augmentcode.com/"
@@ -407,13 +427,15 @@ def create_app(augment_base_url, chat_endpoint, timeout):
             # 决定是否使用流式响应
             if request.stream:
                 return StreamingResponse(
-                    stream_augment_response(augment_base_url, api_key, augment_request, request.model, chat_endpoint, timeout),
+                    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)
-                
+                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(
@@ -459,10 +481,11 @@ def create_app(augment_base_url, chat_endpoint, timeout):
 
     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密钥
@@ -470,7 +493,7 @@ async def handle_sync_request(base_url, api_key, augment_request, model_name, ch
         model_name: 模型名称
         chat_endpoint: 聊天端点
         timeout: 请求超时时间
-        
+
     Returns:
         OpenAI格式的聊天完成响应
     """
@@ -485,7 +508,7 @@ async def handle_sync_request(base_url, api_key, augment_request, model_name, ch
                 "Accept": "*/*"
             }
         )
-        
+
         if response.status_code != 200:
             logger.error(f"Augment API error: {response.status_code} - {response.text}")
             raise HTTPException(
@@ -499,7 +522,7 @@ async def handle_sync_request(base_url, api_key, augment_request, model_name, ch
                     }
                 }
             )
-        
+
         # 处理流式响应，合并为完整响应
         full_response = ""
         for line in response.text.split("\n"):
@@ -510,11 +533,11 @@ async def handle_sync_request(base_url, api_key, augment_request, model_name, ch
                         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()}",
@@ -537,10 +560,11 @@ async def handle_sync_request(base_url, api_key, augment_request, model_name, ch
             )
         )
 
+
 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密钥
@@ -548,22 +572,22 @@ async def stream_augment_response(base_url, api_key, augment_request, model_name
         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": "chrome",
-                    "Accept": "*/*"
-                }
+                    "POST",
+                    f"{base_url.rstrip('/')}/{chat_endpoint}",
+                    json=augment_request.dict(),
+                    headers={
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {api_key}",
+                        "User-Agent": "chrome",
+                        "Accept": "*/*"
+                    }
             ) as response:
 
                 if response.status_code != 200:
@@ -572,11 +596,11 @@ async def stream_augment_response(base_url, api_key, augment_request, model_name
                     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,
@@ -592,19 +616,19 @@ async def stream_augment_response(base_url, api_key, augment_request, model_name
                 )
                 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,
@@ -622,7 +646,7 @@ async def stream_augment_response(base_url, api_key, augment_request, model_name
                             yield f"data: {response_data}\n\n"
                     except json.JSONDecodeError:
                         logger.warning(f"Failed to parse JSON: {line}")
-                
+
                 # 发送完成信号
                 final_response = ChatCompletionStreamResponse(
                     id=chat_id,
@@ -638,10 +662,10 @@ async def stream_augment_response(base_url, api_key, augment_request, model_name
                 )
                 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"
@@ -652,47 +676,48 @@ async def stream_augment_response(base_url, api_key, augment_request, model_name
             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", 
+        "--augment-url",
         default="https://d6.api.augmentcode.com/",
         help="Augment API基础URL"
     )
-    
+
     parser.add_argument(
-        "--chat-endpoint", 
+        "--chat-endpoint",
         default="chat-stream",
         help="Augment聊天端点路径"
     )
-    
+
     parser.add_argument(
-        "--host", 
+        "--host",
         default="0.0.0.0",
         help="服务器主机地址"
     )
-    
+
     parser.add_argument(
-        "--port", 
-        type=int, 
+        "--port",
+        type=int,
         default=8686,
         help="服务器端口"
     )
-    
+
     parser.add_argument(
-        "--timeout", 
-        type=int, 
+        "--timeout",
+        type=int,
         default=120,
         help="API请求超时时间（秒）"
     )
-    
+
     parser.add_argument(
-        "--debug", 
+        "--debug",
         action="store_true",
         help="启用调试模式"
     )
@@ -702,9 +727,10 @@ def parse_args():
         default="d18",
         help="Augment API租户ID (域名前缀)"
     )
-    
+
     return parser.parse_args()
 
+
 #################################################
 # 主程序
 #################################################
@@ -712,11 +738,11 @@ def 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参数
@@ -725,25 +751,26 @@ def main():
     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}:7860")
     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, 
+        app,
+        host=args.host,
         port=7860,
         log_level="info" if not args.debug else "debug"
     )
 
+
 if __name__ == "__main__":
     main()