refac: async pipelines requests

This commit is contained in:
Timothy Jaeryang Baek
2026-02-11 18:24:30 -06:00
parent d02e826c9d
commit 2372b70031
+108 -78
View File
@@ -13,7 +13,6 @@ import aiohttp
import os import os
import logging import logging
import shutil import shutil
import requests
from pydantic import BaseModel from pydantic import BaseModel
from starlette.responses import FileResponse from starlette.responses import FileResponse
from typing import Optional from typing import Optional
@@ -217,7 +216,7 @@ async def upload_pipeline(
os.makedirs(upload_folder, exist_ok=True) os.makedirs(upload_folder, exist_ok=True)
file_path = os.path.join(upload_folder, filename) file_path = os.path.join(upload_folder, filename)
r = None response = None
try: try:
# Save the uploaded file # Save the uploaded file
with open(file_path, "wb") as buffer: with open(file_path, "wb") as buffer:
@@ -226,16 +225,25 @@ async def upload_pipeline(
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
with open(file_path, "rb") as f: headers = {"Authorization": f"Bearer {key}"}
files = {"file": f}
r = requests.post( async with aiohttp.ClientSession(trust_env=True) as session:
f"{url}/pipelines/upload", form_data = aiohttp.FormData()
headers={"Authorization": f"Bearer {key}"}, form_data.add_field(
files=files, "file",
open(file_path, "rb"),
filename=filename,
content_type="application/octet-stream",
) )
r.raise_for_status() async with session.post(
data = r.json() f"{url}/pipelines/upload",
headers=headers,
data=form_data,
ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as response:
response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -244,10 +252,10 @@ async def upload_pipeline(
detail = None detail = None
status_code = status.HTTP_404_NOT_FOUND status_code = status.HTTP_404_NOT_FOUND
if r is not None: if response is not None:
status_code = r.status_code status_code = response.status
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
@@ -272,21 +280,22 @@ class AddPipelineForm(BaseModel):
async def add_pipeline( async def add_pipeline(
request: Request, form_data: AddPipelineForm, user=Depends(get_admin_user) request: Request, form_data: AddPipelineForm, user=Depends(get_admin_user)
): ):
r = None response = None
try: try:
urlIdx = form_data.urlIdx urlIdx = form_data.urlIdx
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
r = requests.post( async with aiohttp.ClientSession(trust_env=True) as session:
f"{url}/pipelines/add", async with session.post(
headers={"Authorization": f"Bearer {key}"}, f"{url}/pipelines/add",
json={"url": form_data.url}, headers={"Authorization": f"Bearer {key}"},
) json={"url": form_data.url},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
r.raise_for_status() ) as response:
data = r.json() response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -294,16 +303,18 @@ async def add_pipeline(
log.exception(f"Connection error: {e}") log.exception(f"Connection error: {e}")
detail = None detail = None
if r is not None: if response is not None:
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
pass pass
raise HTTPException( raise HTTPException(
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), status_code=(
response.status if response is not None else status.HTTP_404_NOT_FOUND
),
detail=detail if detail else "Pipeline not found", detail=detail if detail else "Pipeline not found",
) )
@@ -317,21 +328,22 @@ class DeletePipelineForm(BaseModel):
async def delete_pipeline( async def delete_pipeline(
request: Request, form_data: DeletePipelineForm, user=Depends(get_admin_user) request: Request, form_data: DeletePipelineForm, user=Depends(get_admin_user)
): ):
r = None response = None
try: try:
urlIdx = form_data.urlIdx urlIdx = form_data.urlIdx
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
r = requests.delete( async with aiohttp.ClientSession(trust_env=True) as session:
f"{url}/pipelines/delete", async with session.delete(
headers={"Authorization": f"Bearer {key}"}, f"{url}/pipelines/delete",
json={"id": form_data.id}, headers={"Authorization": f"Bearer {key}"},
) json={"id": form_data.id},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
r.raise_for_status() ) as response:
data = r.json() response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -339,16 +351,18 @@ async def delete_pipeline(
log.exception(f"Connection error: {e}") log.exception(f"Connection error: {e}")
detail = None detail = None
if r is not None: if response is not None:
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
pass pass
raise HTTPException( raise HTTPException(
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), status_code=(
response.status if response is not None else status.HTTP_404_NOT_FOUND
),
detail=detail if detail else "Pipeline not found", detail=detail if detail else "Pipeline not found",
) )
@@ -357,15 +371,19 @@ async def delete_pipeline(
async def get_pipelines( async def get_pipelines(
request: Request, urlIdx: Optional[int] = None, user=Depends(get_admin_user) request: Request, urlIdx: Optional[int] = None, user=Depends(get_admin_user)
): ):
r = None response = None
try: try:
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
r = requests.get(f"{url}/pipelines", headers={"Authorization": f"Bearer {key}"}) async with aiohttp.ClientSession(trust_env=True) as session:
async with session.get(
r.raise_for_status() f"{url}/pipelines",
data = r.json() headers={"Authorization": f"Bearer {key}"},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as response:
response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -373,16 +391,18 @@ async def get_pipelines(
log.exception(f"Connection error: {e}") log.exception(f"Connection error: {e}")
detail = None detail = None
if r is not None: if response is not None:
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
pass pass
raise HTTPException( raise HTTPException(
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), status_code=(
response.status if response is not None else status.HTTP_404_NOT_FOUND
),
detail=detail if detail else "Pipeline not found", detail=detail if detail else "Pipeline not found",
) )
@@ -394,17 +414,19 @@ async def get_pipeline_valves(
pipeline_id: str, pipeline_id: str,
user=Depends(get_admin_user), user=Depends(get_admin_user),
): ):
r = None response = None
try: try:
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
r = requests.get( async with aiohttp.ClientSession(trust_env=True) as session:
f"{url}/{pipeline_id}/valves", headers={"Authorization": f"Bearer {key}"} async with session.get(
) f"{url}/{pipeline_id}/valves",
headers={"Authorization": f"Bearer {key}"},
r.raise_for_status() ssl=AIOHTTP_CLIENT_SESSION_SSL,
data = r.json() ) as response:
response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -412,16 +434,18 @@ async def get_pipeline_valves(
log.exception(f"Connection error: {e}") log.exception(f"Connection error: {e}")
detail = None detail = None
if r is not None: if response is not None:
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
pass pass
raise HTTPException( raise HTTPException(
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), status_code=(
response.status if response is not None else status.HTTP_404_NOT_FOUND
),
detail=detail if detail else "Pipeline not found", detail=detail if detail else "Pipeline not found",
) )
@@ -433,18 +457,19 @@ async def get_pipeline_valves_spec(
pipeline_id: str, pipeline_id: str,
user=Depends(get_admin_user), user=Depends(get_admin_user),
): ):
r = None response = None
try: try:
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
r = requests.get( async with aiohttp.ClientSession(trust_env=True) as session:
f"{url}/{pipeline_id}/valves/spec", async with session.get(
headers={"Authorization": f"Bearer {key}"}, f"{url}/{pipeline_id}/valves/spec",
) headers={"Authorization": f"Bearer {key}"},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
r.raise_for_status() ) as response:
data = r.json() response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -452,16 +477,18 @@ async def get_pipeline_valves_spec(
log.exception(f"Connection error: {e}") log.exception(f"Connection error: {e}")
detail = None detail = None
if r is not None: if response is not None:
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
pass pass
raise HTTPException( raise HTTPException(
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), status_code=(
response.status if response is not None else status.HTTP_404_NOT_FOUND
),
detail=detail if detail else "Pipeline not found", detail=detail if detail else "Pipeline not found",
) )
@@ -474,19 +501,20 @@ async def update_pipeline_valves(
form_data: dict, form_data: dict,
user=Depends(get_admin_user), user=Depends(get_admin_user),
): ):
r = None response = None
try: try:
url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx] url = request.app.state.config.OPENAI_API_BASE_URLS[urlIdx]
key = request.app.state.config.OPENAI_API_KEYS[urlIdx] key = request.app.state.config.OPENAI_API_KEYS[urlIdx]
r = requests.post( async with aiohttp.ClientSession(trust_env=True) as session:
f"{url}/{pipeline_id}/valves/update", async with session.post(
headers={"Authorization": f"Bearer {key}"}, f"{url}/{pipeline_id}/valves/update",
json={**form_data}, headers={"Authorization": f"Bearer {key}"},
) json={**form_data},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
r.raise_for_status() ) as response:
data = r.json() response.raise_for_status()
data = await response.json()
return {**data} return {**data}
except Exception as e: except Exception as e:
@@ -495,15 +523,17 @@ async def update_pipeline_valves(
detail = None detail = None
if r is not None: if response is not None:
try: try:
res = r.json() res = await response.json()
if "detail" in res: if "detail" in res:
detail = res["detail"] detail = res["detail"]
except Exception: except Exception:
pass pass
raise HTTPException( raise HTTPException(
status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND), status_code=(
response.status if response is not None else status.HTTP_404_NOT_FOUND
),
detail=detail if detail else "Pipeline not found", detail=detail if detail else "Pipeline not found",
) )