Adjust the directory structure to prepare for merging into Master.

2023-09-27 11:40:46 -07:00
parent 5146bc1871
commit 030e4c4f52
115 changed files with 413 additions and 160 deletions
@@ -0,0 +1,144 @@
+import asyncio
+
+import openai
+from openai.error import RateLimitError, APIError, Timeout
+
+from jarvis import CFG
+from jarvis.logger import logger
+from typing import Callable
+
+openai.api_key = CFG.openai_api_key
+if CFG.openai_url_base is not None:
+    openai.api_base = CFG.openai_url_base
+
+print_total_cost = CFG.debug_mode
+
+
+async def acreate_chat_completion_once(
+        messages: list,  # type: ignore
+        model: str | None = None,
+        temperature: float = CFG.temperature,
+        max_tokens: int | None = None,
+        deployment_id=None,
+        request_timeout=40,
+        **kwargs
+) -> str:
+    """
+    Create a chat completion and update the cost.
+    Args:
+    messages (list): The list of messages to send to the API.
+    model (str): The model to use for the API call.
+    temperature (float): The temperature to use for the API call.
+    max_tokens (int): The maximum number of tokens for the API call.
+    Returns:
+    str: The AI's response.
+    """
+    if deployment_id is not None:
+        response = await openai.ChatCompletion.acreate(
+            deployment_id=deployment_id,
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            request_timeout=request_timeout,
+            **kwargs
+        )
+    else:
+        response = await openai.ChatCompletion.acreate(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            request_timeout=request_timeout,
+            **kwargs
+        )
+    if CFG.debug_mode:
+        logger.debug(f"Response: {response}")
+    # prompt_tokens = response.usage.prompt_tokens
+    # completion_tokens = response.usage.completion_tokens
+    return response
+
+
+# Overly simple abstraction until we create something better
+# simple retry mechanism when getting a rate error or a bad gateway
+async def acreate_chat_completion(
+        messages: list[dict],
+        model: str = None,
+        temperature: float = CFG.temperature,
+        max_tokens: int = None,
+        request_timeout: int = 40,
+        num_retries=3,
+        on_single_request_timeout: Callable = None,
+        **kwargs
+):
+    """Create a chat completion using the OpenAI API
+
+    Args:
+        messages (List[dict]): The messages to send to the chat completion
+        model (str, optional): The model to use. Defaults to None.
+        temperature (float, optional): The temperature to use. Defaults to 0.9.
+        max_tokens (int, optional): The max tokens to use. Defaults to None.
+        request_timeout (int, optional): The request_timeout of a single openai request.
+        num_retries (int, optional): The max retries.
+        on_single_request_timeout (Callable, optional): This function will be called each time a single openai request
+            timeout, must be an async function, the last timeout will not emit callback.
+
+    Returns:
+        str: The response from the chat completion
+    """
+    if CFG.debug_mode:
+        logger.debug(
+            f"Creating chat completion with model {model}, temperature {temperature}, max_tokens {max_tokens}"
+        )
+
+    response = None
+
+    for attempt in range(num_retries):
+        backoff = min(2 ** (attempt + 2), 8)
+        try:
+            if CFG.use_azure:
+                response = await acreate_chat_completion_once(
+                    deployment_id=CFG.get_azure_deployment_id_for_model(model),
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    request_timeout=request_timeout,
+                    **kwargs
+                )
+            else:
+                response = await acreate_chat_completion_once(
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    request_timeout=request_timeout,
+                    **kwargs
+                )
+            break
+        except RateLimitError:
+            if CFG.debug_mode:
+                logger.debug(f"Error: Reached rate limit, passing...")
+        except (APIError, Timeout) as e:
+            if isinstance(e, Timeout):
+                if on_single_request_timeout:
+                    await on_single_request_timeout(num_retries < num_retries - 1)
+            if e.http_status != 502:
+                raise
+            if attempt == num_retries - 1:
+                raise
+        if CFG.debug_mode:
+            logger.debug(
+                f"Error: API Bad gateway. Waiting {backoff} seconds..."
+            )
+        await asyncio.sleep(backoff)
+    if response is None:
+        logger.error(f"Failed to get response from GPT after {num_retries} retries")
+        raise RuntimeError(f"Failed to get response after {num_retries} retries")
+
+    choice_message = response.choices[0].message
+    content = choice_message.get("content")
+    if content is None:
+        return "function_call", {k: v for k, v in choice_message["function_call"].items()}
+    else:
+        return "content", content