Merge branch 'master' into workspace-handles-v2

daras_ai_v2/base.py

@@ -358,21 +358,7 @@ def render(self):
                     with gui.nav_item(url, active=tab == self.tab):
                         gui.html(tab.title)
 
-            if (
-                self.current_pr
-                and not self.current_pr.is_root()
-                and self.tab == RecipeTabs.run
-            ):
-                with gui.div(
-                    className="container-margin-reset d-none d-md-block",
-                    style=dict(
-                        position="absolute",
-                        top="50%",
-                        right="0",
-                        transform="translateY(-50%)",
-                    ),
-                ):
-                    self._render_saved_timestamp(self.current_pr)
+                    self._render_saved_timestamp()
         with gui.nav_tab_content():
             self.render_selected_tab()
 
@@ -490,9 +476,27 @@ def _render_unpublished_changes_indicator(self):
         ):
             gui.html("Unpublished changes")
 
-    def _render_saved_timestamp(self, pr: PublishedRun):
-        with gui.tag("span", className="text-muted"):
-            gui.write(f"{get_relative_time(pr.updated_at)}")
+    def _render_saved_timestamp(self):
+        sr, pr = self.current_sr_pr
+        if not (pr and self.tab == RecipeTabs.run):
+            return
+        if pr.saved_run_id == sr.id or pr.is_root():
+            dt = pr.updated_at
+        else:
+            dt = sr.updated_at
+        with (
+            gui.div(
+                className="container-margin-reset d-none d-md-block",
+                style=dict(
+                    position="absolute",
+                    top="50%",
+                    right="0",
+                    transform="translateY(-50%)",
+                ),
+            ),
+            gui.tag("span", className="text-muted"),
+        ):
+            gui.write(get_relative_time(dt))
 
     def _render_options_button_with_dialog(self):
         ref = gui.use_alert_dialog(key="options-modal")

daras_ai_v2/bots.py

@@ -21,7 +21,6 @@
     MessageAttachment,
     BotIntegration,
 )
-from daras_ai.image_input import truncate_text_words
 from daras_ai_v2.asr import run_google_translate, should_translate_lang
 from daras_ai_v2.base import BasePage, RecipeRunState, StateKeys
 from daras_ai_v2.csv_lines import csv_encode_row, csv_decode_row
@@ -245,28 +244,21 @@ def translate_response(self, text: str | None) -> str:
             return text or ""
 
 
-def parse_bot_html(
-    text: str | None,
-    max_title_len: int = 20,
-    max_id_len: int = 256,
-) -> tuple[list[ReplyButton], str]:
+def parse_bot_html(text: str | None) -> tuple[list[ReplyButton], str]:
     from pyquery import PyQuery as pq
 
     if not text:
         return [], text
     doc = pq(f"<root>{text}</root>")
     buttons = [
         ReplyButton(
-            id=truncate_text_words(
-                # parsed by _handle_interactive_msg
-                csv_encode_row(
-                    idx + 1,
-                    btn.attrib.get("gui-target") or "input_prompt",
-                    btn.text,
-                ),
-                max_id_len,
+            # parsed by _handle_interactive_msg
+            id=csv_encode_row(
+                idx + 1,
+                btn.attrib.get("gui-target") or "input_prompt",
+                btn.text,
             ),
-            title=truncate_text_words(btn.text, max_title_len),
+            title=btn.text,
         )
         for idx, btn in enumerate(doc("button") or [])
         if btn.text

daras_ai_v2/exceptions.py

@@ -7,6 +7,7 @@
 from loguru import logger
 from requests import HTTPError
 from starlette.status import HTTP_402_PAYMENT_REQUIRED
+from starlette.status import HTTP_401_UNAUTHORIZED
 
 from daras_ai_v2 import settings
 
@@ -103,6 +104,20 @@ def __init__(self, user: "AppUser", sr: "SavedRun"):
         super().__init__(message, status_code=HTTP_402_PAYMENT_REQUIRED)
 
 
+class OneDriveAuth(UserError):
+    def __init__(self, auth_url):
+        message = f"""
+<p>
+OneDrive access is currently unavailable. 
+</p>
+
+<p>
+<a href="{auth_url}">LOGIN</a> to your OneDrive account to enable access to your files.
+</p>
+"""
+        super().__init__(message, status_code=HTTP_401_UNAUTHORIZED)
+
+
 FFMPEG_ERR_MSG = (
     "Unsupported File Format\n\n"
     "We encountered an issue processing your file as it appears to be in a format not supported by our system or may be corrupted. "

daras_ai_v2/facebook_bots.py

@@ -2,7 +2,11 @@
 from furl import furl
 
 from bots.models import BotIntegration, Platform, Conversation
-from daras_ai.image_input import upload_file_from_bytes, get_mimetype_from_response
+from daras_ai.image_input import (
+    upload_file_from_bytes,
+    get_mimetype_from_response,
+    truncate_text_words,
+)
 from daras_ai.text_format import markdown_to_wa
 from daras_ai_v2 import settings
 from daras_ai_v2.asr import (
@@ -281,24 +285,38 @@ def retrieve_wa_media_by_id(
     return content, media_info["mime_type"]
 
 
-def _build_msg_buttons(buttons: list[ReplyButton], msg: dict) -> list[dict]:
+def _build_msg_buttons(
+    buttons: list[ReplyButton],
+    msg: dict,
+    *,
+    max_title_len: int = 20,
+    max_id_len: int = 256,
+) -> list[dict]:
     ret = []
     for i in range(0, len(buttons), 3):
+        buttons = [
+            {
+                "type": "reply",
+                "reply": {
+                    "id": truncate_text_words(
+                        button["id"],
+                        max_id_len,
+                    ),
+                    "title": truncate_text_words(
+                        button["title"],
+                        max_title_len,
+                    ),
+                },
+            }
+            for button in buttons[i : i + 3]
+        ]
         ret.append(
             {
                 "type": "interactive",
                 "interactive": {
                     "type": "button",
                     **msg,
-                    "action": {
-                        "buttons": [
-                            {
-                                "type": "reply",
-                                "reply": {"id": button["id"], "title": button["title"]},
-                            }
-                            for button in buttons[i : i + 3]
-                        ],
-                    },
+                    "action": {"buttons": buttons},
                 },
             }
         )

daras_ai_v2/functional.py

@@ -31,7 +31,9 @@ def apply_parallel(
     if not max_workers:
         yield
         return []
-    with ThreadPoolExecutor(max_workers=max_workers) as pool:
+    with ThreadPoolExecutor(
+        max_workers=max_workers, initializer=get_initializer()
+    ) as pool:
         fs = [pool.submit(fn, *args) for args in zip(*iterables)]
         length = len(fs)
         ret = [None] * length
@@ -52,7 +54,9 @@ def fetch_parallel(
     max_workers = max_workers or max(map(len, iterables))
     if not max_workers:
         return
-    with ThreadPoolExecutor(max_workers=max_workers) as pool:
+    with ThreadPoolExecutor(
+        max_workers=max_workers, initializer=get_initializer()
+    ) as pool:
         fs = [pool.submit(fn, *args) for args in zip(*iterables)]
         for fut in as_completed(fs):
             yield fut.result()
@@ -90,7 +94,9 @@ def map_parallel(
     max_workers = max_workers or max(map(len, iterables))
     if not max_workers:
         return []
-    with ThreadPoolExecutor(max_workers=max_workers) as pool:
+    with ThreadPoolExecutor(
+        max_workers=max_workers, initializer=get_initializer()
+    ) as pool:
         return list(pool.map(fn, *iterables))
 
 
@@ -104,7 +110,9 @@ def map_parallel_ascompleted(
     max_workers = max_workers or max(map(len, iterables))
     if not max_workers:
         return
-    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+    with ThreadPoolExecutor(
+        max_workers=max_workers, initializer=get_initializer()
+    ) as executor:
         fs = {
             executor.submit(fn, *args): arg_id
             for arg_id, args in zip(ids, zip(*iterables))
@@ -124,3 +132,14 @@ def shape(seq: typing.Sequence | np.ndarray) -> tuple[int, ...]:
         return (len(seq),) + shape(seq[0])
     else:
         return ()
+
+
+def get_initializer() -> typing.Callable:
+    from celeryapp.tasks import threadlocal
+
+    parent = threadlocal.__dict__
+
+    def initializer():
+        threadlocal.__dict__.update(parent)
+
+    return initializer

daras_ai_v2/language_model.py

@@ -80,7 +80,7 @@ class LargeLanguageModels(Enum):
     # https://platform.openai.com/docs/models#o3-mini
     o3_mini = LLMSpec(
         label="o3-mini (openai)",
-        model_id="o3-mini-2025-01-31",
+        model_id=("openai-o3-mini-prod-eastus2-1", "o3-mini-2025-01-31"),
         llm_api=LLMApis.openai,
         context_window=200_000,
         price=13,
@@ -92,7 +92,7 @@ class LargeLanguageModels(Enum):
     # https://platform.openai.com/docs/models#o1
     o1 = LLMSpec(
         label="o1 (openai)",
-        model_id="o1-2024-12-17",
+        model_id=("openai-o1-prod-eastus2-1", "o1-2024-12-17"),
         llm_api=LLMApis.openai,
         context_window=200_000,
         price=50,
@@ -1395,14 +1395,14 @@ def get_openai_client(model: str):
         client = openai.AzureOpenAI(
             api_key=settings.AZURE_OPENAI_KEY_CA,
             azure_endpoint=settings.AZURE_OPENAI_ENDPOINT_CA,
-            api_version="2023-10-01-preview",
+            api_version="2024-12-01-preview",
             max_retries=0,
         )
     elif model.startswith(AZURE_OPENAI_MODEL_PREFIX) and "-eastus2-" in model:
         client = openai.AzureOpenAI(
             api_key=settings.AZURE_OPENAI_KEY_EASTUS2,
             azure_endpoint=settings.AZURE_OPENAI_ENDPOINT_EASTUS2,
-            api_version="2023-10-01-preview",
+            api_version="2024-12-01-preview",
             max_retries=0,
         )
     else:

daras_ai_v2/onedrive_downloader.py

@@ -0,0 +1,100 @@
+import base64
+
+import requests
+from furl import furl
+
+from bots.models import SavedRun
+from daras_ai_v2.exceptions import UserError
+from daras_ai_v2.exceptions import raise_for_status, OneDriveAuth
+from routers.onedrive_api import (
+    generate_onedrive_auth_url,
+    get_access_token_from_refresh_token,
+)
+
+
+def is_onedrive_url(f: furl) -> bool:
+    if f.host == "1drv.ms":
+        return True
+    elif f.host == "onedrive.live.com":
+        raise UserError(
+            "Direct onedrive.live.com links are not supported. Please provide a shareable OneDrive link (from Share > Copy Link) E.g. https://1drv.ms/xxx"
+        )
+
+
+_url_encode_translation = str.maketrans({"/": "_", "+": "-", "=": ""})
+
+
+def encode_onedrive_url(sharing_url: str) -> str:
+    # https://learn.microsoft.com/en-us/onedrive/developer/rest-api/api/shares_get
+    base64_value = base64.b64encode(sharing_url.encode()).decode()
+    encoded_url = base64_value.translate(_url_encode_translation)
+    return f"u!{encoded_url}"
+
+
+def onedrive_download(mime_type: str, export_links: dict):
+    download_url = export_links.get(mime_type)
+    if not download_url:
+        raise ValueError(
+            "Download URL not found in export_links. Cannot download file."
+        )
+    r = requests.get(download_url)
+    raise_for_status(r)
+    file_content = r.content
+    return file_content, mime_type
+
+
+def onedrive_meta(f_url: str, sr: SavedRun, *, try_refresh: bool = True):
+    # check if saved run workspace has onedrive_access_token and onedrive_refresh_token
+    if not (
+        sr.workspace
+        and sr.workspace.onedrive_access_token
+        and sr.workspace.onedrive_refresh_token
+    ):
+        raise OneDriveAuth(generate_onedrive_auth_url(sr.id))
+    try:
+        encoded_url = encode_onedrive_url(f_url)
+        headers = {"Authorization": f"Bearer {sr.workspace.onedrive_access_token}"}
+        r = requests.get(
+            f"https://graph.microsoft.com/v1.0/shares/{encoded_url}/driveItem",
+            headers=headers,
+        )
+        raise_for_status(r)
+        metadata = r.json()
+
+        if "folder" in metadata:
+            raise UserError(
+                "Folders & OneNote files are not supported yet. Please remove them from your Knowledge base."
+            )
+
+        return metadata
+
+    except requests.HTTPError as e:
+        if e.response.status_code == 401 and try_refresh:
+            try:
+                (
+                    sr.workspace.onedrive_access_token,
+                    sr.workspace.onedrive_refresh_token,
+                ) = get_access_token_from_refresh_token(
+                    sr.workspace.onedrive_refresh_token
+                )
+                sr.workspace.save(update_fields=["onedrive_access_token"])
+                return onedrive_meta(f_url, sr, try_refresh=False)
+            except requests.HTTPError:
+                raise OneDriveAuth(generate_onedrive_auth_url(sr.id))
+
+        elif e.response.status_code == 403:
+            raise UserError(
+                message=f"""
+<p>
+<a href="{f_url}" target="_blank">This document </a> is not accessible by "{sr.workspace.onedrive_user_name}".
+Please share the document with this account (Share > Manage Access).
+</p>
+<p>
+Alternatively, <a href="{generate_onedrive_auth_url(sr.id)}" target="_blank">Login</a> with a OneDrive account that can access this file.
+Note that you can only be logged in to one OneDrive account at a time.
+</p>
+"""
+            )
+
+        else:
+            raise

daras_ai_v2/settings.py

@@ -396,6 +396,10 @@
 SLACK_CLIENT_ID = config("SLACK_CLIENT_ID", "")
 SLACK_CLIENT_SECRET = config("SLACK_CLIENT_SECRET", "")
 
+ONEDRIVE_CLIENT_ID = config("ONEDRIVE_CLIENT_ID", "")
+ONEDRIVE_CLIENT_SECRET = config("ONEDRIVE_CLIENT_SECRET", "")
+
+
 TALK_JS_APP_ID = config("TALK_JS_APP_ID", "")
 TALK_JS_SECRET_KEY = config("TALK_JS_SECRET_KEY", "")