add redis lock timeout to prevent infinite deadlocks

use loguru for logging
update llm pricing in doc search
add azure openai fallbacks using aifail
add gpt-4-16k

daras_ai_v2/redis_cache.py

@@ -9,6 +9,9 @@
 from daras_ai_v2 import settings
 
 
+LOCK_TIMEOUT_SEC = 10 * 60
+
+
 @lru_cache
 def get_redis_cache():
     return redis.Redis.from_url(settings.REDIS_CACHE_URL)
@@ -27,7 +30,14 @@ def wrapper(*args, **kwargs):
         # get the redis cache
         redis_cache = get_redis_cache()
         # lock the cache key so that only one thread can run the function
-        with redis_cache.lock(os.path.join(cache_key, "lock")):
+        lock = redis_cache.lock(
+            name=os.path.join(cache_key, "lock"), timeout=LOCK_TIMEOUT_SEC
+        )
+        try:
+            lock.acquire()
+        except redis.exceptions.LockError:
+            pass
+        try:
             cache_val = redis_cache.get(cache_key)
             # if the cache exists, return it
             if cache_val:
@@ -38,5 +48,10 @@ def wrapper(*args, **kwargs):
                 cache_val = pickle.dumps(result)
                 redis_cache.set(cache_key, cache_val)
                 return result
+        finally:
+            try:
+                lock.release()
+            except redis.exceptions.LockError:
+                pass
 
     return wrapper

daras_ai_v2/settings.py

@@ -293,6 +293,9 @@
 AZURE_IMAGE_MODERATION_ENDPOINT = config("AZURE_IMAGE_MODERATION_ENDPOINT", "")
 AZURE_IMAGE_MODERATION_KEY = config("AZURE_IMAGE_MODERATION_KEY", "")
 
+AZURE_OPENAI_ENDPOINT = config("AZURE_OPENAI_ENDPOINT", "")
+AZURE_OPENAI_KEY = config("AZURE_OPENAI_KEY", "")
+
 DEEPGRAM_API_KEY = config("DEEPGRAM_API_KEY", "")
 
 ELEVEN_LABS_API_KEY = config("ELEVEN_LABS_API_KEY", "")

daras_ai_v2/vector_search.py

@@ -401,7 +401,7 @@ def get_embeds_for_doc(
         ]
     # get doc embeds in batches
     embeds = []
-    batch_size = 100
+    batch_size = 16  # azure openai limits
     texts = [m["title"] + " | " + m["snippet"] for m in metas]
     for i in range(0, len(texts), batch_size):
         # progress = int(i / len(texts) * 100)

gooey_ui/pubsub.py

@@ -6,6 +6,7 @@
 
 import redis
 from fastapi.encoders import jsonable_encoder
+from loguru import logger
 
 from daras_ai_v2 import settings
 
@@ -41,7 +42,7 @@ def realtime_push(channel: str, value: typing.Any = "ping"):
     msg = json.dumps(jsonable_encoder(value))
     r.set(channel, msg)
     r.publish(channel, json.dumps(time()))
-    print(f"publish {channel!r}")
+    logger.info(f"publish {channel=}")
 
 
 # def use_state(

pyproject.toml

@@ -80,6 +80,8 @@ pyyaml = "^6.0.1"
 ua-parser = "^0.18.0"
 user-agents = "^2.2.0"
 openpyxl = "^3.1.2"
+loguru = "^0.7.2"
+aifail = "^0.1.0"
 
 [tool.poetry.group.dev.dependencies]
 watchdog = "^2.1.9"

recipes/CompareLLM.py

@@ -8,7 +8,11 @@
 from bots.models import Workflow
 from daras_ai_v2.base import BasePage
 from daras_ai_v2.enum_selector_widget import enum_multiselect
-from daras_ai_v2.language_model import run_language_model, LargeLanguageModels
+from daras_ai_v2.language_model import (
+    run_language_model,
+    LargeLanguageModels,
+    llm_price,
+)
 from daras_ai_v2.language_model_settings_widgets import language_model_settings
 from daras_ai_v2.loom_video_widget import youtube_video
 from daras_ai_v2.prompt_vars import prompt_vars_widget, render_prompt_vars
@@ -131,27 +135,10 @@ def get_raw_price(self, state: dict) -> int:
         selected_models = state.get("selected_models", [])
         total = 0
         for name in selected_models:
-            match name:
-                case LargeLanguageModels.gpt_4.name:
-                    total += 10
-                case LargeLanguageModels.gpt_3_5_turbo_16k.name:
-                    total += 2
-                case LargeLanguageModels.gpt_3_5_turbo.name:
-                    total += 1
-                case LargeLanguageModels.text_davinci_003.name | LargeLanguageModels.code_davinci_002.name:
-                    total += 10
-                case LargeLanguageModels.text_curie_001.name:
-                    total += 5
-                case LargeLanguageModels.text_babbage_001.name:
-                    total += 2
-                case LargeLanguageModels.text_ada_001.name:
-                    total += 1
-                case LargeLanguageModels.palm2_text.name:
-                    total += 15
-                case LargeLanguageModels.palm2_chat.name:
-                    total += 10
-                case LargeLanguageModels.llama2_70b_chat.name:
-                    total += 5
+            try:
+                total += llm_price[LargeLanguageModels[name]]
+            except KeyError:
+                total += 5
         return total * state.get("num_outputs", 1)
 
     def related_workflows(self) -> list:

recipes/DocSearch.py

@@ -1,7 +1,5 @@
-import datetime
 import typing
 
-import jinja2
 from furl import furl
 from pydantic import BaseModel
 
@@ -16,7 +14,7 @@
 from daras_ai_v2.language_model import (
     run_language_model,
     LargeLanguageModels,
-    model_max_tokens,
+    llm_price,
 )
 from daras_ai_v2.language_model_settings_widgets import language_model_settings
 from daras_ai_v2.loom_video_widget import youtube_video
@@ -201,13 +199,10 @@ def run_v2(
 
     def get_raw_price(self, state: dict) -> float:
         name = state.get("selected_model")
-        match name:
-            case LargeLanguageModels.gpt_4.name:
-                return 60
-            case LargeLanguageModels.gpt_3_5_turbo_16k.name:
-                return 20
-            case _:
-                return 10
+        try:
+            return llm_price[LargeLanguageModels[name]] * 2
+        except KeyError:
+            return 10
 
 
 def render_documents(state, label="**Documents**", *, key="documents"):