Merge pull request WecoAI#21 from dexhunter/fix/beta_models

🐛 fix model issues with beta limitation

aide/backend/__init__.py

@@ -33,6 +33,14 @@ def query(
         "max_tokens": max_tokens,
     }
 
+    # Handle models with beta limitations
+    # ref: https://platform.openai.com/docs/guides/reasoning/beta-limitations
+    if model.startswith("o1-"):
+        if system_message:
+            user_message = system_message
+        system_message = None
+        model_kwargs["temperature"] = 1
+
     query_func = backend_anthropic.query if "claude-" in model else backend_openai.query
     output, req_time, in_tok_count, out_tok_count, info = query_func(
         system_message=compile_prompt_to_md(system_message) if system_message else None,

aide/backend/backend_anthropic.py

@@ -2,23 +2,25 @@
 
 import time
 
-from anthropic import Anthropic, RateLimitError
-from .utils import FunctionSpec, OutputType, opt_messages_to_list
-from funcy import notnone, once, retry, select_values
+from .utils import FunctionSpec, OutputType, opt_messages_to_list, backoff_create
+from funcy import notnone, once, select_values
+import anthropic
 
-_client: Anthropic = None  # type: ignore
+_client: anthropic.Anthropic = None  # type: ignore
 
-RATELIMIT_RETRIES = 5
-retry_exp = retry(RATELIMIT_RETRIES, errors=RateLimitError, timeout=lambda a: 2 ** (a + 1))  # type: ignore
+ANTHROPIC_TIMEOUT_EXCEPTIONS = (
+    anthropic.RateLimitError,
+    anthropic.APIConnectionError,
+    anthropic.APITimeoutError,
+    anthropic.InternalServerError,
+)
 
 
 @once
 def _setup_anthropic_client():
     global _client
-    _client = Anthropic()
+    _client = anthropic.Anthropic(max_retries=0)
 
-
-@retry_exp
 def query(
     system_message: str | None,
     user_message: str | None,
@@ -48,7 +50,12 @@ def query(
     messages = opt_messages_to_list(None, user_message)
 
     t0 = time.time()
-    message = _client.messages.create(messages=messages, **filtered_kwargs)  # type: ignore
+    message = backoff_create(
+        _client.messages.create,
+        ANTHROPIC_TIMEOUT_EXCEPTIONS,
+        messages=messages,
+        **filtered_kwargs,
+    )
     req_time = time.time() - t0
 
     assert len(message.content) == 1 and message.content[0].type == "text"

aide/backend/backend_openai.py

@@ -4,25 +4,26 @@
 import logging
 import time
 
-from .utils import FunctionSpec, OutputType, opt_messages_to_list
-from funcy import notnone, once, retry, select_values
-from openai import OpenAI, RateLimitError
+from .utils import FunctionSpec, OutputType, opt_messages_to_list, backoff_create
+from funcy import notnone, once, select_values
+import openai
 
 logger = logging.getLogger("aide")
 
-_client: OpenAI = None  # type: ignore
-
-RATELIMIT_RETRIES = 5
-retry_exp = retry(RATELIMIT_RETRIES, errors=RateLimitError, timeout=lambda a: 2 ** (a + 1))  # type: ignore
+_client: openai.OpenAI = None  # type: ignore
 
+OPENAI_TIMEOUT_EXCEPTIONS = (
+    openai.RateLimitError,
+    openai.APIConnectionError,
+    openai.APITimeoutError,
+    openai.InternalServerError,
+)
 
 @once
 def _setup_openai_client():
     global _client
-    _client = OpenAI(max_retries=3)
-
+    _client = openai.OpenAI(max_retries=0)
 
-@retry_exp
 def query(
     system_message: str | None,
     user_message: str | None,
@@ -40,7 +41,12 @@ def query(
         filtered_kwargs["tool_choice"] = func_spec.openai_tool_choice_dict
 
     t0 = time.time()
-    completion = _client.chat.completions.create(messages=messages, **filtered_kwargs)  # type: ignore
+    completion = backoff_create(
+        _client.chat.completions.create,
+        OPENAI_TIMEOUT_EXCEPTIONS,
+        messages=messages,
+        **filtered_kwargs,
+    )
     req_time = time.time() - t0
 
     choice = completion.choices[0]

aide/backend/utils.py

@@ -8,6 +8,27 @@
 OutputType = str | FunctionCallType
 
 
+import backoff
+import logging
+from typing import Callable
+
+logger = logging.getLogger("aide")
+
+
+@backoff.on_predicate(
+    wait_gen=backoff.expo,
+    max_value=60,
+    factor=1.5,
+)
+def backoff_create(
+    create_fn: Callable, retry_exceptions: list[Exception], *args, **kwargs
+):
+    try:
+        return create_fn(*args, **kwargs)
+    except retry_exceptions as e:
+        logger.info(f"Backoff exception: {e}")
+        return False
+
 def opt_messages_to_list(
     system_message: str | None, user_message: str | None
 ) -> list[dict[str, str]]:

requirements.txt

@@ -88,4 +88,5 @@ pdf2image
 PyPDF
 pyocr
 pyarrow
-xlrd
+xlrd
+backoff