add full async support for newapi

This commit is contained in:
2025-07-23 19:51:20 +08:00
parent 1c05f934fe
commit d97813ec1a

View File

@@ -8,7 +8,7 @@ author_url: https://zhuangyumin.dev
""" """
import re import re
import requests import aiohttp
import json import json
import time import time
import tiktoken import tiktoken
@@ -150,26 +150,29 @@ class Pipe:
"cost": 0.0 "cost": 0.0
} }
def pipes(self) -> List[dict]: async def pipes(self) -> List[dict]:
"""Fetch available models from OpenRouter API""" """Fetch available models from NewAPI asynchronously"""
if not self.valves.NEWAPI_API_KEY: if not self.valves.NEWAPI_API_KEY:
return [{"id": "error", "name": "API Key not provided"}] return [{"id": "error", "name": "API Key not provided"}]
try: try:
headers = {"Authorization": f"Bearer {self.valves.NEWAPI_API_KEY}"} headers = {"Authorization": f"Bearer {self.valves.NEWAPI_API_KEY}"}
response = requests.get(
f"{self.valves.NEWAPI_BASE_URL}/models", headers=headers
)
if response.status_code != 200: async with aiohttp.ClientSession() as session:
async with session.get(
f"{self.valves.NEWAPI_BASE_URL}/models",
headers=headers,
timeout=aiohttp.ClientTimeout(total=30)
) as response:
if response.status != 200:
return [ return [
{ {
"id": "error", "id": "error",
"name": f"Error fetching models: {response.status_code}", "name": f"Error fetching models: {response.status}",
} }
] ]
models_data = response.json() models_data = await response.json()
# Extract model information # Extract model information
models = [] models = []
@@ -193,7 +196,7 @@ class Pipe:
return [{"id": "error", "name": f"Error: {str(e)}"}] return [{"id": "error", "name": f"Error: {str(e)}"}]
async def _report_api_call_direct(self, usage_info: dict, user_email: str, model_id: str, __event_emitter__: Callable[[Any], Awaitable[None]]): async def _report_api_call_direct(self, usage_info: dict, user_email: str, model_id: str, __event_emitter__: Callable[[Any], Awaitable[None]]):
"""Report API call to upstream reporting service using direct usage information""" """Report API call to upstream reporting service using direct usage information asynchronously"""
if not self.valves.REPORT_API_URL or not self.valves.REPORT_API_KEY: if not self.valves.REPORT_API_URL or not self.valves.REPORT_API_KEY:
return return
@@ -214,7 +217,7 @@ class Pipe:
"cost_usd": cost_usd "cost_usd": cost_usd
} }
# Send to reporting API # Send to reporting API asynchronously
headers = { headers = {
"Authorization": f"Bearer {self.valves.REPORT_API_KEY}", "Authorization": f"Bearer {self.valves.REPORT_API_KEY}",
"Content-Type": "application/json" "Content-Type": "application/json"
@@ -222,17 +225,17 @@ class Pipe:
report_url = f"{self.valves.REPORT_API_URL.rstrip('/')}/api/record_api_call" report_url = f"{self.valves.REPORT_API_URL.rstrip('/')}/api/record_api_call"
response = requests.post( async with aiohttp.ClientSession() as session:
async with session.post(
report_url, report_url,
headers=headers, headers=headers,
json=api_call_record, json=api_call_record,
timeout=30 timeout=aiohttp.ClientTimeout(total=30)
) ) as response:
if response.status == 200:
if response.status_code == 200:
print(f"Successfully reported API call for user {user_email}") print(f"Successfully reported API call for user {user_email}")
else: else:
print(f"Failed to report API call: {response.status_code}") print(f"Failed to report API call: {response.status}")
info = f"input: {input_tokens} | output: {output_tokens} | cost: {cost_usd:.6f}" info = f"input: {input_tokens} | output: {output_tokens} | cost: {cost_usd:.6f}"
await __event_emitter__( await __event_emitter__(
@@ -300,7 +303,7 @@ class Pipe:
return self.stream_response(url, headers, payload, user_email, model_id, __event_emitter__, model_name) return self.stream_response(url, headers, payload, user_email, model_id, __event_emitter__, model_name)
else: else:
return await self.non_stream_response(url, headers, payload, user_email, model_id, __event_emitter__, model_name) return await self.non_stream_response(url, headers, payload, user_email, model_id, __event_emitter__, model_name)
except requests.exceptions.RequestException as e: except aiohttp.ClientError as e:
print(f"Request failed: {e}") print(f"Request failed: {e}")
return f"Error: Request failed: {e}" return f"Error: Request failed: {e}"
except Exception as e: except Exception as e:
@@ -311,14 +314,20 @@ class Pipe:
"""Handle non-streaming responses and wrap reasoning in <think> tags if present""" """Handle non-streaming responses and wrap reasoning in <think> tags if present"""
try: try:
print( print(
f"Sending non-streaming request to OpenRouter: {json.dumps(payload)[:200]}..." f"Sending non-streaming request to NewAPI: {json.dumps(payload)[:200]}..."
) )
response = requests.post(url, headers=headers, json=payload, timeout=90)
if response.status_code != 200: async with aiohttp.ClientSession() as session:
error_message = f"HTTP Error {response.status_code}" async with session.post(
url,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=90)
) as response:
if response.status != 200:
error_message = f"HTTP Error {response.status}"
try: try:
error_data = response.json() error_data = await response.json()
print(f"Error response: {json.dumps(error_data)}") print(f"Error response: {json.dumps(error_data)}")
if "error" in error_data: if "error" in error_data:
if ( if (
@@ -330,14 +339,15 @@ class Pipe:
error_message += f": {error_data['error']}" error_message += f": {error_data['error']}"
except Exception as e: except Exception as e:
print(f"Failed to parse error response: {e}") print(f"Failed to parse error response: {e}")
error_message += f": {response.text[:500]}" error_text = await response.text()
error_message += f": {error_text[:500]}"
# Log request payload for debugging # Log request payload for debugging
print(f"Request that caused error: {json.dumps(payload)}") print(f"Request that caused error: {json.dumps(payload)}")
raise Exception(error_message) raise Exception(error_message)
res = response.json() res = await response.json()
print(f"OpenRouter response keys: {list(res.keys())}") print(f"NewAPI response keys: {list(res.keys())}")
# Check if we have choices in the response # Check if we have choices in the response
if not res.get("choices") or len(res["choices"]) == 0: if not res.get("choices") or len(res["choices"]) == 0:
@@ -384,14 +394,17 @@ class Pipe:
async def stream_response(self, url, headers, payload, user_email, model_id, __event_emitter__: Callable[[Any], Awaitable[None]], model_name: str): async def stream_response(self, url, headers, payload, user_email, model_id, __event_emitter__: Callable[[Any], Awaitable[None]], model_name: str):
"""Stream reasoning tokens in real-time with proper tag management""" """Stream reasoning tokens in real-time with proper tag management"""
try: try:
response = requests.post( async with aiohttp.ClientSession() as session:
url, headers=headers, json=payload, stream=True, timeout=90 async with session.post(
) url,
headers=headers,
if response.status_code != 200: json=payload,
error_message = f"HTTP Error {response.status_code}" timeout=aiohttp.ClientTimeout(total=90)
) as response:
if response.status != 200:
error_message = f"HTTP Error {response.status}"
try: try:
error_data = response.json() error_data = await response.json()
error_message += ( error_message += (
f": {error_data.get('error', {}).get('message', '')}" f": {error_data.get('error', {}).get('message', '')}"
) )
@@ -405,15 +418,21 @@ class Pipe:
accumulated_content = "" # Accumulate all content for token calculation accumulated_content = "" # Accumulate all content for token calculation
accumulated_reasoning = "" # Accumulate all reasoning for token calculation accumulated_reasoning = "" # Accumulate all reasoning for token calculation
# Process the response stream # Process the response stream asynchronously
for line in response.iter_lines(): async for line_bytes in response.content:
if not line: if not line_bytes:
continue continue
line_text = line.decode("utf-8") line_text = line_bytes.decode("utf-8").strip()
if not line_text.startswith("data: "):
# Handle multiple lines in a single chunk
for line in line_text.split('\n'):
if not line.strip():
continue continue
elif line_text == "data: [DONE]":
if not line.startswith("data: "):
continue
elif line == "data: [DONE]":
# Handle citations at the end # Handle citations at the end
if latest_citations: if latest_citations:
citation_list = [f"1. {l}" for l in latest_citations] citation_list = [f"1. {l}" for l in latest_citations]
@@ -443,7 +462,7 @@ class Pipe:
break break
try: try:
chunk = json.loads(line_text[6:]) chunk = json.loads(line[6:])
if "choices" in chunk and chunk["choices"]: if "choices" in chunk and chunk["choices"]:
choice = chunk["choices"][0] choice = chunk["choices"][0]