chore: format

This commit is contained in:
Timothy Jaeryang Baek
2025-02-09 22:20:47 -08:00
parent 426f8f29ad
commit 60095598ec
57 changed files with 387 additions and 42 deletions

View File

@@ -120,18 +120,12 @@ class OpenSearchClient:
return None
query_body = {
"query": {
"bool": {
"filter": []
}
},
"query": {"bool": {"filter": []}},
"_source": ["text", "metadata"],
}
for field, value in filter.items():
query_body["query"]["bool"]["filter"].append({
"term": {field: value}
})
query_body["query"]["bool"]["filter"].append({"term": {field: value}})
size = limit if limit else 10
@@ -139,7 +133,7 @@ class OpenSearchClient:
result = self.client.search(
index=f"{self.index_prefix}_{collection_name}",
body=query_body,
size=size
size=size,
)
return self._result_to_get_result(result)

View File

@@ -25,13 +25,10 @@ def search_jina(api_key: str, query: str, count: int) -> list[SearchResult]:
"Accept": "application/json",
"Content-Type": "application/json",
"Authorization": api_key,
"X-Retain-Images": "none"
"X-Retain-Images": "none",
}
payload = {
"q": query,
"count": count if count <= 10 else 10
}
payload = {"q": query, "count": count if count <= 10 else 10}
url = str(URL(jina_search_endpoint))
response = requests.post(url, headers=headers, json=payload)

View File

@@ -560,10 +560,14 @@ def transcribe(request: Request, file_path):
# Extract transcript from Deepgram response
try:
transcript = response_data["results"]["channels"][0]["alternatives"][0].get("transcript", "")
transcript = response_data["results"]["channels"][0]["alternatives"][
0
].get("transcript", "")
except (KeyError, IndexError) as e:
log.error(f"Malformed response from Deepgram: {str(e)}")
raise Exception("Failed to parse Deepgram response - unexpected response format")
raise Exception(
"Failed to parse Deepgram response - unexpected response format"
)
data = {"text": transcript.strip()}
# Save transcript

View File

@@ -1424,11 +1424,11 @@ async def upload_model(
os.makedirs(UPLOAD_DIR, exist_ok=True)
# --- P1: save file locally ---
chunk_size = 1024 * 1024 * 2 # 2 MB chunks
chunk_size = 1024 * 1024 * 2 # 2 MB chunks
with open(file_path, "wb") as out_f:
while True:
chunk = file.file.read(chunk_size)
#log.info(f"Chunk: {str(chunk)}") # DEBUG
# log.info(f"Chunk: {str(chunk)}") # DEBUG
if not chunk:
break
out_f.write(chunk)
@@ -1436,15 +1436,15 @@ async def upload_model(
async def file_process_stream():
nonlocal ollama_url
total_size = os.path.getsize(file_path)
log.info(f"Total Model Size: {str(total_size)}") # DEBUG
log.info(f"Total Model Size: {str(total_size)}") # DEBUG
# --- P2: SSE progress + calculate sha256 hash ---
file_hash = calculate_sha256(file_path, chunk_size)
log.info(f"Model Hash: {str(file_hash)}") # DEBUG
log.info(f"Model Hash: {str(file_hash)}") # DEBUG
try:
with open(file_path, "rb") as f:
bytes_read = 0
while chunk := f.read(chunk_size):
while chunk := f.read(chunk_size):
bytes_read += len(chunk)
progress = round(bytes_read / total_size * 100, 2)
data_msg = {
@@ -1460,25 +1460,23 @@ async def upload_model(
response = requests.post(url, data=f)
if response.ok:
log.info(f"Uploaded to /api/blobs") # DEBUG
log.info(f"Uploaded to /api/blobs") # DEBUG
# Remove local file
os.remove(file_path)
# Create model in ollama
model_name, ext = os.path.splitext(file.filename)
log.info(f"Created Model: {model_name}") # DEBUG
log.info(f"Created Model: {model_name}") # DEBUG
create_payload = {
"model": model_name,
# Reference the file by its original name => the uploaded blob's digest
"files": {
file.filename: f"sha256:{file_hash}"
},
"files": {file.filename: f"sha256:{file_hash}"},
}
log.info(f"Model Payload: {create_payload}") # DEBUG
log.info(f"Model Payload: {create_payload}") # DEBUG
# Call ollama /api/create
#https://github.com/ollama/ollama/blob/main/docs/api.md#create-a-model
# https://github.com/ollama/ollama/blob/main/docs/api.md#create-a-model
create_resp = requests.post(
url=f"{ollama_url}/api/create",
headers={"Content-Type": "application/json"},
@@ -1486,7 +1484,7 @@ async def upload_model(
)
if create_resp.ok:
log.info(f"API SUCCESS!") # DEBUG
log.info(f"API SUCCESS!") # DEBUG
done_msg = {
"done": True,
"blob": f"sha256:{file_hash}",
@@ -1506,4 +1504,4 @@ async def upload_model(
res = {"error": str(e)}
yield f"data: {json.dumps(res)}\n\n"
return StreamingResponse(file_process_stream(), media_type="text/event-stream")
return StreamingResponse(file_process_stream(), media_type="text/event-stream")

View File

@@ -94,7 +94,7 @@ class S3StorageProvider(StorageProvider):
aws_secret_access_key=S3_SECRET_ACCESS_KEY,
)
self.bucket_name = S3_BUCKET_NAME
self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
"""Handles uploading of the file to S3 storage."""
@@ -108,7 +108,7 @@ class S3StorageProvider(StorageProvider):
)
except ClientError as e:
raise RuntimeError(f"Error uploading file to S3: {e}")
def get_file(self, file_path: str) -> str:
"""Handles downloading of the file from S3 storage."""
try:
@@ -137,7 +137,8 @@ class S3StorageProvider(StorageProvider):
if "Contents" in response:
for content in response["Contents"]:
# Skip objects that were not uploaded from open-webui in the first place
if not content["Key"].startswith(self.key_prefix): continue
if not content["Key"].startswith(self.key_prefix):
continue
self.s3_client.delete_object(
Bucket=self.bucket_name, Key=content["Key"]
@@ -150,11 +151,12 @@ class S3StorageProvider(StorageProvider):
# The s3 key is the name assigned to an object. It excludes the bucket name, but includes the internal path and the file name.
def _extract_s3_key(self, full_file_path: str) -> str:
return '/'.join(full_file_path.split("//")[1].split("/")[1:])
return "/".join(full_file_path.split("//")[1].split("/")[1:])
def _get_local_file_path(self, s3_key: str) -> str:
return f"{UPLOAD_DIR}/{s3_key.split('/')[-1]}"
class GCSStorageProvider(StorageProvider):
def __init__(self):
self.bucket_name = GCS_BUCKET_NAME

View File

@@ -245,7 +245,7 @@ def get_gravatar_url(email):
def calculate_sha256(file_path, chunk_size):
#Compute SHA-256 hash of a file efficiently in chunks
# Compute SHA-256 hash of a file efficiently in chunks
sha256 = hashlib.sha256()
with open(file_path, "rb") as f:
while chunk := f.read(chunk_size):

View File

@@ -142,13 +142,17 @@ class OAuthManager:
log.debug(f"Oauth Groups claim: {oauth_claim}")
log.debug(f"User oauth groups: {user_oauth_groups}")
log.debug(f"User's current groups: {[g.name for g in user_current_groups]}")
log.debug(f"All groups available in OpenWebUI: {[g.name for g in all_available_groups]}")
log.debug(
f"All groups available in OpenWebUI: {[g.name for g in all_available_groups]}"
)
# Remove groups that user is no longer a part of
for group_model in user_current_groups:
if group_model.name not in user_oauth_groups:
# Remove group from user
log.debug(f"Removing user from group {group_model.name} as it is no longer in their oauth groups")
log.debug(
f"Removing user from group {group_model.name} as it is no longer in their oauth groups"
)
user_ids = group_model.user_ids
user_ids = [i for i in user_ids if i != user.id]
@@ -174,7 +178,9 @@ class OAuthManager:
gm.name == group_model.name for gm in user_current_groups
):
# Add user to group
log.debug(f"Adding user to group {group_model.name} as it was found in their oauth groups")
log.debug(
f"Adding user to group {group_model.name} as it was found in their oauth groups"
)
user_ids = group_model.user_ids
user_ids.append(user.id)
@@ -289,7 +295,9 @@ class OAuthManager:
base64_encoded_picture = base64.b64encode(
picture
).decode("utf-8")
guessed_mime_type = mimetypes.guess_type(picture_url)[0]
guessed_mime_type = mimetypes.guess_type(
picture_url
)[0]
if guessed_mime_type is None:
# assume JPG, browsers are tolerant enough of image formats
guessed_mime_type = "image/jpeg"

View File

@@ -12,6 +12,7 @@ from fpdf import FPDF
from open_webui.env import STATIC_DIR, FONTS_DIR
from open_webui.models.chats import ChatTitleMessagesForm
class PDFGenerator:
"""
Description: