add full async support for newapi
This commit is contained in:
@@ -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]
|
||||||
|
Reference in New Issue
Block a user