MilkyBloomVibeCode/patch_llm.py at main · QuangVoAI/MilkyBloomVibeCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import re

with open("agentic-ai/python/agents/llm_client.py", "r") as f:
    code = f.read()

# 1. Update imports
new_imports = """import random
from config import (
    GROQ_API_KEY,
    GROQ_API_KEYS,
    GROQ_BASE_URL,
    GROQ_MODEL,
    GROQ_MODEL_FAST,
    GROQ_MODEL_SMART,
    GROQ_HTTP_REFERER,
    GROQ_X_TITLE,
    FEATHERLESS_API_KEY,
    FEATHERLESS_BASE_URL,
    FEATHERLESS_MODEL,
    FEATHERLESS_MODEL_FAST,
    FEATHERLESS_MODEL_SMART,
    FEATHERLESS_HTTP_REFERER,
    FEATHERLESS_X_TITLE,
)"""

code = re.sub(r'from config import \([^)]+\)', new_imports, code)

# 2. Add Groq candidate functions
groq_functions = """
_GROQ_KEY_COOLDOWNS = {}

def _cooldown_groq_key(api_key: str, error: Exception | str | None) -> float:
    retry_after = _extract_retry_after(error)
    seconds = retry_after if retry_after is not None else 60.0
    seconds = max(10.0, min(float(seconds), 3600.0))
    _GROQ_KEY_COOLDOWNS[api_key] = time.monotonic() + seconds
    return seconds

def _is_groq_key_available(api_key: str) -> bool:
    expires_at = _GROQ_KEY_COOLDOWNS.get(api_key, 0)
    return expires_at <= time.monotonic()

def _groq_api_key_candidates() -> list[str]:
    keys = _unique_nonempty(list(GROQ_API_KEYS or []))
    if not keys and GROQ_API_KEY:
        keys = [GROQ_API_KEY.strip()]
    return keys

def _available_groq_api_key_candidates() -> list[tuple[int, str]]:
    keys = _groq_api_key_candidates()
    candidates = [
        (index, key)
        for index, key in enumerate(keys, start=1)
        if _is_groq_key_available(key)
    ]
    random.shuffle(candidates)
    return candidates

def _groq_model_candidates(requested_model: str) -> list[str]:
    return _unique_nonempty([
        requested_model,
        GROQ_MODEL_SMART,
        GROQ_MODEL_FAST,
        GROQ_MODEL,
    ])

"""
code = code.replace("def _cooldown_featherless_key", groq_functions + "def _cooldown_featherless_key")

# 3. Patch featherless_chat_complete to use Groq for the main loop
old_chat_func = """    model_candidates = _featherless_model_candidates(model)
    key_candidates = _available_featherless_api_key_candidates()
    last_error: Exception | None = None
    abort_groq = False

    if _featherless_api_key_candidates() and not key_candidates:
        last_error = RuntimeError("All Groq API keys are temporarily rate-limited.")"""

new_chat_func = """    model_candidates = _groq_model_candidates(model)
    key_candidates = _available_groq_api_key_candidates()
    last_error: Exception | None = None
    abort_groq = False

    if _groq_api_key_candidates() and not key_candidates:
        last_error = RuntimeError("All Groq API keys are temporarily rate-limited.")"""

code = code.replace(old_chat_func, new_chat_func)

code = code.replace("base_url=FEATHERLESS_BASE_URL,", "base_url=GROQ_BASE_URL,")
code = code.replace("referer=FEATHERLESS_HTTP_REFERER,", "referer=GROQ_HTTP_REFERER,")
code = code.replace("title=FEATHERLESS_X_TITLE,", "title=GROQ_X_TITLE,")
# Oops, the replace above replaces too much. We only want to replace in the groq try block.
# I will use a more precise regex.