moved utility to common and updated data_url parsing logic

This commit is contained in:
Hardik Shah 2024-09-12 11:58:04 -07:00
parent 5f49dce839
commit 487e16dc3f
4 changed files with 60 additions and 30 deletions

View file

@ -8,9 +8,6 @@
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import base64
import mimetypes
import os
from typing import List, Optional, Protocol
from llama_models.schema_utils import json_schema_type, webmethod
@ -29,21 +26,6 @@ class MemoryBankDocument(BaseModel):
metadata: Dict[str, Any] = Field(default_factory=dict)
def data_url_from_file(file_path: str) -> URL:
if not os.path.exists(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
with open(file_path, "rb") as file:
file_content = file.read()
base64_content = base64.b64encode(file_content).decode("utf-8")
mime_type, _ = mimetypes.guess_type(file_path)
data_url = f"data:{mime_type};base64,{base64_content}"
return URL(uri=data_url)
@json_schema_type
class MemoryBankType(Enum):
vector = "vector"