import os import uuid import json import time import asyncio import random from curl_cffi.requests import AsyncSession from fastapi import FastAPI, Request, HTTPException, Depends, status from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from fastapi.responses import StreamingResponse from dotenv import load_dotenv import secrets from pydantic import BaseModel, Field from typing import List, Optional, Dict, Any, Literal, Union # Load environment variables from .env file load_dotenv() # --- 并发请求配置 --- CONCURRENT_REQUESTS = 1 # 可自定义并发请求数量 # --- 重试配置 --- MAX_RETRIES = 3 RETRY_DELAY = 1 # 秒 # --- Models (Integrated from models.py) --- # Input Models (OpenAI-like) class ChatMessage(BaseModel): role: Literal["system", "user", "assistant"] content: str class ChatCompletionRequest(BaseModel): messages: List[ChatMessage] model: str = "notion-proxy" # Model name can be passed, but we map to Notion's model stream: bool = False # Add other potential OpenAI params if needed, though they might not map directly # max_tokens: Optional[int] = None # temperature: Optional[float] = None # space_id and thread_id are now handled globally via environment variables notion_model: str = "anthropic-opus-4" # Default Notion model, can be overridden # Notion Models class NotionTranscriptConfigValue(BaseModel): type: str = "markdown-chat" model: str # e.g., "anthropic-opus-4" class NotionTranscriptItem(BaseModel): type: Literal["config", "user", "markdown-chat"] value: Union[List[List[str]], str, NotionTranscriptConfigValue] class NotionDebugOverrides(BaseModel): cachedInferences: Dict = Field(default_factory=dict) annotationInferences: Dict = Field(default_factory=dict) emitInferences: bool = False class NotionRequestBody(BaseModel): traceId: str = Field(default_factory=lambda: str(uuid.uuid4())) spaceId: str transcript: List[NotionTranscriptItem] # threadId is removed, createThread will be set to true createThread: bool = True debugOverrides: NotionDebugOverrides = Field(default_factory=NotionDebugOverrides) generateTitle: bool = False saveAllThreadOperations: bool = True # Output Models (OpenAI SSE) class ChoiceDelta(BaseModel): content: Optional[str] = None class Choice(BaseModel): index: int = 0 delta: ChoiceDelta finish_reason: Optional[Literal["stop", "length"]] = None class ChatCompletionChunk(BaseModel): id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4()}") object: str = "chat.completion.chunk" created: int = Field(default_factory=lambda: int(time.time())) model: str = "notion-proxy" # Or could reflect the underlying Notion model choices: List[Choice] # Models for /v1/models Endpoint class Model(BaseModel): id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "notion" # Or specify based on actual model origin if needed class ModelList(BaseModel): object: str = "list" data: List[Model] # --- Configuration --- NOTION_API_URL = "https://www.notion.so/api/v3/runInferenceTranscript" # IMPORTANT: Load the Notion cookie securely from environment variables NOTION_COOKIE = os.getenv("NOTION_COOKIE") NOTION_SPACE_ID = os.getenv("NOTION_SPACE_ID") if not NOTION_COOKIE: print("Error: NOTION_COOKIE environment variable not set.") # Consider raising HTTPException or exiting in a real app if not NOTION_SPACE_ID: print("Warning: NOTION_SPACE_ID environment variable not set. Using a default UUID.") # Using a default might not be ideal, depends on Notion's behavior # Consider raising an error instead: raise ValueError("NOTION_SPACE_ID not set") NOTION_SPACE_ID = str(uuid.uuid4()) # Default or raise error # --- Authentication --- EXPECTED_TOKEN = os.getenv("PROXY_AUTH_TOKEN", "default_token") # Default token security = HTTPBearer() def authenticate(credentials: HTTPAuthorizationCredentials = Depends(security)): """Compares provided token with the expected token.""" correct_token = secrets.compare_digest(credentials.credentials, EXPECTED_TOKEN) if not correct_token: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authentication credentials", # WWW-Authenticate header removed for Bearer ) return True # Indicate successful authentication # --- FastAPI App --- app = FastAPI() # --- Helper Functions --- def build_notion_request(request_data: ChatCompletionRequest) -> NotionRequestBody: """Transforms OpenAI-style messages to Notion transcript format.""" transcript = [ NotionTranscriptItem( type="config", value=NotionTranscriptConfigValue(model=request_data.notion_model) ) ] for message in request_data.messages: # Map 'assistant' role to 'markdown-chat', all others to 'user' if message.role == "assistant": # Notion uses "markdown-chat" for assistant replies in the transcript history transcript.append(NotionTranscriptItem(type="markdown-chat", value=message.content)) else: # Map user, system, and any other potential roles to 'user' transcript.append(NotionTranscriptItem(type="user", value=[[message.content]])) # Use globally configured spaceId, set createThread=True return NotionRequestBody( spaceId=NOTION_SPACE_ID, # From environment variable transcript=transcript, createThread=True, # Always create a new thread # Generate a new traceId for each request traceId=str(uuid.uuid4()), # Explicitly set debugOverrides, generateTitle, and saveAllThreadOperations debugOverrides=NotionDebugOverrides( cachedInferences={}, annotationInferences={}, emitInferences=False ), generateTitle=False, saveAllThreadOperations=False ) async def check_first_response_line(session: AsyncSession, notion_request_body: NotionRequestBody, headers: dict, request_id: int): """检查响应的第一行，判断是否为500错误""" try: # 当并发请求数大于1时，添加随机延迟以避免同时到达 if CONCURRENT_REQUESTS > 1: delay = random.uniform(0, 1.0) print(f"并发请求 {request_id} 延迟 {delay:.2f}秒") await asyncio.sleep(delay) # 为每个并发请求创建独立的请求体，生成新的traceId request_body_copy = notion_request_body.model_copy() request_body_copy.traceId = str(uuid.uuid4()) response = await session.post( NOTION_API_URL, json=request_body_copy.model_dump(), headers=headers, stream=True ) if response.status_code != 200: return None, response, f"HTTP {response.status_code}" # 读取第一行来检查是否是错误 buffer = "" async for chunk in response.aiter_content(): if isinstance(chunk, bytes): chunk = chunk.decode('utf-8') buffer += chunk # 尝试解析第一个完整的JSON行 lines = buffer.split('\n') for line in lines: line = line.strip() if line: try: data = json.loads(line) if (data.get("type") == "error" and data.get("message") and "error code 500" in data.get("message", "")): print(f"并发请求 {request_id} 检测到500错误: {data}") return None, response, "500 error" else: # 正常响应，返回response和已读取的buffer print(f"并发请求 {request_id} 响应正常") return (response, buffer), None, None except json.JSONDecodeError: continue return None, response, "No valid response" except Exception as e: print(f"并发请求 {request_id} 发生异常: {e}") return None, None, str(e) async def stream_notion_response_single(session: AsyncSession, response, initial_buffer: str, chunk_id: str, created_time: int): """处理单个响应的流式输出""" buffer = initial_buffer # 首先处理已经读取的buffer中的内容 lines = buffer.split('\n') buffer = lines[-1] for line in lines[:-1]: line = line.strip() if not line: continue try: data = json.loads(line) if data.get("type") == "markdown-chat" and isinstance(data.get("value"), str): content_chunk = data["value"] if content_chunk: chunk_obj = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=content_chunk))] ) yield f"data: {chunk_obj.model_dump_json()}\n\n" elif "recordMap" in data: print("Detected recordMap, stopping stream.") # 继续处理剩余的buffer if buffer.strip(): try: last_data = json.loads(buffer.strip()) if last_data.get("type") == "markdown-chat" and isinstance(last_data.get("value"), str): if last_data["value"]: last_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=last_data["value"]))] ) yield f"data: {last_chunk.model_dump_json()}\n\n" except: pass return except json.JSONDecodeError as e: print(f"Warning: Could not decode JSON line: {line[:100]}... Error: {str(e)}") except Exception as e: print(f"Error processing line: {str(e)}") # 继续读取剩余的响应 async for chunk in response.aiter_content(): if isinstance(chunk, bytes): chunk = chunk.decode('utf-8') buffer += chunk lines = buffer.split('\n') buffer = lines[-1] for line in lines[:-1]: line = line.strip() if not line: continue try: data = json.loads(line) if data.get("type") == "markdown-chat" and isinstance(data.get("value"), str): content_chunk = data["value"] if content_chunk: chunk_obj = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=content_chunk))] ) yield f"data: {chunk_obj.model_dump_json()}\n\n" elif "recordMap" in data: print("Detected recordMap, stopping stream.") if buffer.strip(): try: last_data = json.loads(buffer.strip()) if last_data.get("type") == "markdown-chat" and isinstance(last_data.get("value"), str): if last_data["value"]: last_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=last_data["value"]))] ) yield f"data: {last_chunk.model_dump_json()}\n\n" except: pass return except json.JSONDecodeError as e: print(f"Warning: Could not decode JSON line: {line[:100]}... Error: {str(e)}") except Exception as e: print(f"Error processing line: {str(e)}") async def stream_notion_response(notion_request_body: NotionRequestBody): """Streams the request to Notion and yields OpenAI-compatible SSE chunks.""" # curl_cffi will automatically handle most headers like a real browser # We only need to set specific headers that are necessary headers = { 'accept': 'application/x-ndjson', 'accept-encoding': 'gzip, deflate, br, zstd', 'accept-language': 'en-US,zh;q=0.9', 'content-type': 'application/json', 'dnt': '1', 'notion-audit-log-platform': 'web', 'notion-client-version': '23.13.0.3661', 'origin': 'https://www.notion.so', 'referer': 'https://www.notion.so/', 'priority': 'u=1, i', 'sec-ch-ua-mobile': '?0', 'sec-ch-ua-platform': '"Windows"', 'sec-fetch-dest': 'empty', 'sec-fetch-mode': 'cors', 'sec-fetch-site': 'same-origin', 'cookie': NOTION_COOKIE, 'x-notion-space-id': NOTION_SPACE_ID } # Conditionally add the active user header notion_active_user = os.getenv("NOTION_ACTIVE_USER_HEADER") if notion_active_user: # Checks for None and empty string implicitly headers['x-notion-active-user-header'] = notion_active_user chunk_id = f"chatcmpl-{uuid.uuid4()}" created_time = int(time.time()) # 使用全局重试配置 max_retries = MAX_RETRIES retry_delay = RETRY_DELAY # 首先尝试并发请求 print(f"同时发起 {CONCURRENT_REQUESTS} 个并发请求...") async with AsyncSession(impersonate="chrome136") as session: # 同时创建并发任务（每个都是独立的异步任务） tasks = [] for i in range(CONCURRENT_REQUESTS): task = asyncio.create_task( check_first_response_line(session, notion_request_body, headers, i + 1) ) tasks.append(task) # 等待所有任务完成或找到第一个成功的响应 successful_response = None failed_count = 0 completed_tasks = set() while len(completed_tasks) < CONCURRENT_REQUESTS and not successful_response: # 等待任意一个任务完成 done, pending = await asyncio.wait( [t for t in tasks if t not in completed_tasks], return_when=asyncio.FIRST_COMPLETED ) for task in done: completed_tasks.add(task) result, response, error = await task if result: # 找到成功的响应，立即使用 successful_response = result print(f"找到成功的并发响应，立即使用") # 取消其他还在运行的任务 for t in tasks: if t not in completed_tasks: t.cancel() break else: # 记录失败 failed_count += 1 if error: print(f"并发请求失败: {error}") # 如果有成功的响应，使用它进行流式传输 if successful_response: response, initial_buffer = successful_response print("使用成功的并发响应进行流式传输") # 流式输出响应 async for data in stream_notion_response_single(session, response, initial_buffer, chunk_id, created_time): yield data # Send the final chunk indicating stop final_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(), finish_reason="stop")] ) yield f"data: {final_chunk.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return # 只有当所有并发请求都失败时，才进入重试流程 print(f"所有 {CONCURRENT_REQUESTS} 个并发请求都失败，开始单请求重试流程...") # 进入原有的重试逻辑（不使用并发） for attempt in range(max_retries): try: # Using curl_cffi with chrome136 impersonation for better anti-bot bypass async with AsyncSession(impersonate="chrome136") as session: # Stream the response response = await session.post( NOTION_API_URL, json=notion_request_body.model_dump(), headers=headers, stream=True ) if response.status_code != 200: error_content = await response.atext() print(f"Error from Notion API: {response.status_code}") print(f"Response: {error_content}") raise HTTPException(status_code=response.status_code, detail=f"Notion API Error: {error_content}") # Process streaming response # curl_cffi streaming works differently - we need to read the content in chunks buffer = "" first_line_checked = False is_error_response = False async for chunk in response.aiter_content(): # Decode chunk if it's bytes if isinstance(chunk, bytes): chunk = chunk.decode('utf-8') buffer += chunk # Split by newlines and process complete lines lines = buffer.split('\n') # Keep the last incomplete line in the buffer buffer = lines[-1] for line in lines[:-1]: line = line.strip() if not line: continue try: data = json.loads(line) # 检查第一行是否是500错误响应 if not first_line_checked: first_line_checked = True if (data.get("type") == "error" and data.get("message") and "error code 500" in data.get("message", "")): print(f"检测到Notion API 500错误 (重试 {attempt + 1}/{max_retries}): {data}") is_error_response = True break # 如果不是错误响应，实时流式转发 # Check if it's the type of message containing text chunks if data.get("type") == "markdown-chat" and isinstance(data.get("value"), str): content_chunk = data["value"] if content_chunk: # Only send if there's content chunk_obj = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=content_chunk))] ) yield f"data: {chunk_obj.model_dump_json()}\n\n" # Add logic here to detect the end of the stream if Notion has a specific marker # For now, we assume markdown-chat stops when the main content is done. # If we see a recordMap, it's definitely past the text stream. elif "recordMap" in data: print("Detected recordMap, stopping stream.") # Process any remaining buffer if buffer.strip(): try: last_data = json.loads(buffer.strip()) if last_data.get("type") == "markdown-chat" and isinstance(last_data.get("value"), str): if last_data["value"]: last_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=last_data["value"]))] ) yield f"data: {last_chunk.model_dump_json()}\n\n" except: pass # Exit the loop break except json.JSONDecodeError as e: print(f"Warning: Could not decode JSON line: {line[:100]}... Error: {str(e)}") except Exception as e: print(f"Error processing line: {str(e)}") # Continue processing other lines if is_error_response: break # 如果检测到错误，进行重试 if is_error_response: if attempt < max_retries - 1: print(f"等待 {retry_delay} 秒后重试...") await asyncio.sleep(retry_delay) continue # 重试 else: # 所有重试都失败了，通过流式响应返回错误信息 print("所有重试都失败，返回500错误给客户端") error_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content="Error: Notion API returned error code 500 after all retries"), finish_reason="stop")] ) yield f"data: {error_chunk.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return # 如果没有错误，发送最终的停止信号 # Send the final chunk indicating stop final_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(), finish_reason="stop")] ) yield f"data: {final_chunk.model_dump_json()}\n\n" yield "data: [DONE]\n\n" # 成功完成，退出重试循环 break except HTTPException: # 在流式响应中不能抛出HTTPException，通过流式响应返回错误 if attempt < max_retries - 1: print(f"HTTP异常，等待 {retry_delay} 秒后重试...") await asyncio.sleep(retry_delay) continue else: print("HTTP异常且无更多重试，返回错误信息") error_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content="Error: HTTP exception occurred after all retries"), finish_reason="stop")] ) yield f"data: {error_chunk.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return except Exception as e: print(f"Unexpected error during streaming (attempt {attempt + 1}/{max_retries}): {e}") if attempt < max_retries - 1: print(f"等待 {retry_delay} 秒后重试...") await asyncio.sleep(retry_delay) continue else: print("意外错误且无更多重试，返回错误信息") error_chunk = ChatCompletionChunk( id=chunk_id, created=created_time, choices=[Choice(delta=ChoiceDelta(content=f"Error: Internal server error during streaming: {e}"), finish_reason="stop")] ) yield f"data: {error_chunk.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return # --- API Endpoints --- @app.get("/v1/models", response_model=ModelList) async def list_models(authenticated: bool = Depends(authenticate)): """ Endpoint to list available Notion models, mimicking OpenAI's /v1/models. """ available_models = [ "openai-gpt-4.1", "anthropic-opus-4", "anthropic-sonnet-4" ] model_list = [ Model(id=model_id, owned_by="notion") # created uses default_factory for model_id in available_models ] return ModelList(data=model_list) @app.post("/v1/chat/completions") async def chat_completions(request_data: ChatCompletionRequest, request: Request, authenticated: bool = Depends(authenticate)): """ Endpoint to mimic OpenAI's chat completions, proxying to Notion. """ if not NOTION_COOKIE: raise HTTPException(status_code=500, detail="Server configuration error: Notion cookie not set.") notion_request_body = build_notion_request(request_data) if request_data.stream: return StreamingResponse( stream_notion_response(notion_request_body), media_type="text/event-stream" ) else: # --- Non-Streaming Logic (Optional - Collects stream internally) --- # Note: The primary goal is streaming, but a non-streaming version # might be useful for testing or simpler clients. # This requires collecting all chunks from the async generator. full_response_content = "" final_finish_reason = None chunk_id = f"chatcmpl-{uuid.uuid4()}" # Generate ID for the non-streamed response created_time = int(time.time()) try: async for line in stream_notion_response(notion_request_body): if line.startswith("data: ") and "[DONE]" not in line: try: data_json = line[len("data: "):].strip() if data_json: chunk_data = json.loads(data_json) if chunk_data.get("choices"): delta = chunk_data["choices"][0].get("delta", {}) content = delta.get("content") if content: full_response_content += content finish_reason = chunk_data["choices"][0].get("finish_reason") if finish_reason: final_finish_reason = finish_reason except json.JSONDecodeError: print(f"Warning: Could not decode JSON line in non-streaming mode: {line}") # Construct the final OpenAI-compatible non-streaming response return { "id": chunk_id, "object": "chat.completion", "created": created_time, "model": request_data.model, # Return the model requested by the client "choices": [ { "index": 0, "message": { "role": "assistant", "content": full_response_content, }, "finish_reason": final_finish_reason or "stop", # Default to stop if not explicitly set } ], "usage": { # Note: Token usage is not available from Notion "prompt_tokens": None, "completion_tokens": None, "total_tokens": None, }, } except HTTPException as e: # Re-raise HTTP exceptions from the streaming function raise e except Exception as e: print(f"Error during non-streaming processing: {e}") raise HTTPException(status_code=500, detail="Internal server error processing Notion response") if __name__ == "__main__": import uvicorn print("Starting server. Access at http://localhost:7860") print("Ensure NOTION_COOKIE is set in your .env file or environment.") uvicorn.run(app, host="0.0.0.0", port=7860)