Merge branch 'main' into couchbase-vectordb-support

autogen/agentchat/contrib/gpt_assistant_agent.py

@@ -209,10 +209,12 @@ def _invoke_assistant(
         for message in pending_messages:
             if message["content"].strip() == "":
                 continue
+            # Convert message roles to 'user' or 'assistant', by calling _map_role_for_api, to comply with OpenAI API spec
+            api_role = self._map_role_for_api(message["role"])
             self._openai_client.beta.threads.messages.create(
                 thread_id=assistant_thread.id,
                 content=message["content"],
-                role=message["role"],
+                role=api_role,
             )
 
         # Create a new run to get responses from the assistant
@@ -240,6 +242,28 @@ def _invoke_assistant(
         self._unread_index[sender] = len(self._oai_messages[sender]) + 1
         return True, response
 
+    def _map_role_for_api(self, role: str) -> str:
+        """
+        Maps internal message roles to the roles expected by the OpenAI Assistant API.
+
+        Args:
+            role (str): The role from the internal message.
+
+        Returns:
+            str: The mapped role suitable for the API.
+        """
+        if role in ["function", "tool"]:
+            return "assistant"
+        elif role == "system":
+            return "system"
+        elif role == "user":
+            return "user"
+        elif role == "assistant":
+            return "assistant"
+        else:
+            # Default to 'assistant' for any other roles not recognized by the API
+            return "assistant"
+
     def _get_run_response(self, thread, run):
         """
         Waits for and processes the response of a run from the OpenAI assistant.

autogen/agentchat/conversable_agent.py

@@ -1658,8 +1658,8 @@ async def a_generate_function_call_reply(
         if messages is None:
             messages = self._oai_messages[sender]
         message = messages[-1]
-        if "function_call" in message:
-            func_call = message["function_call"]
+        func_call = message.get("function_call")
+        if func_call:
             func_name = func_call.get("name", "")
             func = self._function_map.get(func_name, None)
             if func and inspect.iscoroutinefunction(func):

autogen/agentchat/groupchat.py

@@ -1398,13 +1398,12 @@ async def a_resume(
             if not message_speaker_agent and message["name"] == self.name:
                 message_speaker_agent = self
 
-            # Add previous messages to each agent (except their own messages and the last message, as we'll kick off the conversation with it)
+            # Add previous messages to each agent (except the last message, as we'll kick off the conversation with it)
             if i != len(messages) - 1:
                 for agent in self._groupchat.agents:
-                    if agent.name != message["name"]:
-                        await self.a_send(
-                            message, self._groupchat.agent_by_name(agent.name), request_reply=False, silent=True
-                        )
+                    await self.a_send(
+                        message, self._groupchat.agent_by_name(agent.name), request_reply=False, silent=True
+                    )
 
                 # Add previous message to the new groupchat, if it's an admin message the name may not match so add the message directly
                 if message_speaker_agent:

autogen/coding/jupyter/docker_jupyter_server.py

@@ -8,7 +8,7 @@
 import uuid
 from pathlib import Path
 from types import TracebackType
-from typing import Dict, Optional, Type, Union
+from typing import Any, Dict, Optional, Type, Union
 
 import docker
 
@@ -59,6 +59,7 @@ def __init__(
         stop_container: bool = True,
         docker_env: Dict[str, str] = {},
         token: Union[str, GenerateToken] = GenerateToken(),
+        **docker_kwargs: Any,
     ):
         """Start a Jupyter kernel gateway server in a Docker container.
 
@@ -77,6 +78,7 @@ def __init__(
             token (Union[str, GenerateToken], optional): Token to use for authentication.
                 If GenerateToken is used, a random token will be generated. Empty string
                 will be unauthenticated.
+            docker_kwargs (Any): Additional keyword arguments to pass to the docker container.
         """
         if container_name is None:
             container_name = f"autogen-jupyterkernelgateway-{uuid.uuid4()}"
@@ -118,6 +120,7 @@ def __init__(
             environment=env,
             publish_all_ports=True,
             name=container_name,
+            **docker_kwargs,
         )
         _wait_for_ready(container)
         container_ports = container.ports

autogen/coding/jupyter/jupyter_client.py

@@ -39,6 +39,10 @@ def _get_headers(self) -> Dict[str, str]:
             return {}
         return {"Authorization": f"token {self._connection_info.token}"}
 
+    def _get_cookies(self) -> str:
+        cookies = self._session.cookies.get_dict()
+        return "; ".join([f"{name}={value}" for name, value in cookies.items()])
+
     def _get_api_base_url(self) -> str:
         protocol = "https" if self._connection_info.use_https else "http"
         port = f":{self._connection_info.port}" if self._connection_info.port else ""
@@ -87,7 +91,7 @@ def restart_kernel(self, kernel_id: str) -> None:
 
     def get_kernel_client(self, kernel_id: str) -> JupyterKernelClient:
         ws_url = f"{self._get_ws_base_url()}/api/kernels/{kernel_id}/channels"
-        ws = websocket.create_connection(ws_url, header=self._get_headers())
+        ws = websocket.create_connection(ws_url, header=self._get_headers(), cookie=self._get_cookies())
         return JupyterKernelClient(ws)
 
 

autogen/logger/file_logger.py

@@ -90,7 +90,7 @@ def log_chat_completion(
         thread_id = threading.get_ident()
         source_name = None
         if isinstance(source, str):
-            source_name = source
+            source_name = getattr(source, "name", "unknown")
         else:
             source_name = source.name
         try:

autogen/oai/client.py

@@ -12,10 +12,12 @@
 from autogen.cache import Cache
 from autogen.io.base import IOStream
 from autogen.logger.logger_utils import get_current_ts
-from autogen.oai.openai_utils import OAI_PRICE1K, get_key, is_valid_api_key
+from autogen.oai.openai_utils import OAI_PRICE1K, get_key
 from autogen.runtime_logging import log_chat_completion, log_new_client, log_new_wrapper, logging_enabled
 from autogen.token_count_utils import count_token
 
+from .rate_limiters import RateLimiter, TimeRateLimiter
+
 TOOL_ENABLED = False
 try:
     import openai
@@ -163,11 +165,7 @@ class OpenAIClient:
 
     def __init__(self, client: Union[OpenAI, AzureOpenAI]):
         self._oai_client = client
-        if (
-            not isinstance(client, openai.AzureOpenAI)
-            and str(client.base_url).startswith(OPEN_API_BASE_URL_PREFIX)
-            and not is_valid_api_key(self._oai_client.api_key)
-        ):
+        if not isinstance(client, openai.AzureOpenAI) and str(client.base_url).startswith(OPEN_API_BASE_URL_PREFIX):
             logger.warning(
                 "The API key specified is not a valid OpenAI format; it won't work with the OpenAI-hosted model."
             )
@@ -207,7 +205,9 @@ def create(self, params: Dict[str, Any]) -> ChatCompletion:
         """
         iostream = IOStream.get_default()
 
-        completions: Completions = self._oai_client.chat.completions if "messages" in params else self._oai_client.completions  # type: ignore [attr-defined]
+        completions: Completions = (
+            self._oai_client.chat.completions if "messages" in params else self._oai_client.completions
+        )  # type: ignore [attr-defined]
         # If streaming is enabled and has messages, then iterate over the chunks of the response.
         if params.get("stream", False) and "messages" in params:
             response_contents = [""] * params.get("n", 1)
@@ -279,7 +279,12 @@ def create(self, params: Dict[str, Any]) -> ChatCompletion:
 
             # Prepare the final ChatCompletion object based on the accumulated data
             model = chunk.model.replace("gpt-35", "gpt-3.5")  # hack for Azure API
-            prompt_tokens = count_token(params["messages"], model)
+            try:
+                prompt_tokens = count_token(params["messages"], model)
+            except NotImplementedError as e:
+                # Catch token calculation error if streaming with customized models.
+                logger.warning(str(e))
+                prompt_tokens = 0
             response = ChatCompletion(
                 id=chunk.id,
                 model=chunk.model,
@@ -422,8 +427,11 @@ def __init__(self, *, config_list: Optional[List[Dict[str, Any]]] = None, **base
 
         self._clients: List[ModelClient] = []
         self._config_list: List[Dict[str, Any]] = []
+        self._rate_limiters: List[Optional[RateLimiter]] = []
 
         if config_list:
+            self._initialize_rate_limiters(config_list)
+
             config_list = [config.copy() for config in config_list]  # make a copy before modifying
             for config in config_list:
                 self._register_default_client(config, openai_config)  # could modify the config
@@ -744,6 +752,7 @@ def yes_or_no_filter(context, response):
                             return response
                         continue  # filter is not passed; try the next config
             try:
+                self._throttle_api_calls(i)
                 request_ts = get_current_ts()
                 response = client.create(params)
             except APITimeoutError as err:
@@ -1037,3 +1046,20 @@ def extract_text_or_completion_object(
             A list of text, or a list of ChatCompletion objects if function_call/tool_calls are present.
         """
         return response.message_retrieval_function(response)
+
+    def _throttle_api_calls(self, idx: int) -> None:
+        """Rate limit api calls."""
+        if self._rate_limiters[idx]:
+            limiter = self._rate_limiters[idx]
+
+            assert limiter is not None
+            limiter.sleep()
+
+    def _initialize_rate_limiters(self, config_list: List[Dict[str, Any]]) -> None:
+        for config in config_list:
+            # Instantiate the rate limiter
+            if "api_rate_limit" in config:
+                self._rate_limiters.append(TimeRateLimiter(config["api_rate_limit"]))
+                del config["api_rate_limit"]
+            else:
+                self._rate_limiters.append(None)

autogen/oai/cohere.py

@@ -148,7 +148,6 @@ def create(self, params: Dict) -> ChatCompletion:
         client_name = params.get("client_name") or "autogen-cohere"
         # Parse parameters to the Cohere API's parameters
         cohere_params = self.parse_params(params)
-
         # Convert AutoGen messages to Cohere messages
         cohere_messages, preamble, final_message = oai_messages_to_cohere_messages(messages, params, cohere_params)
 
@@ -169,13 +168,15 @@ def create(self, params: Dict) -> ChatCompletion:
         cohere_finish = ""
 
         max_retries = 5
+
         for attempt in range(max_retries):
             ans = None
             try:
                 if streaming:
                     response = client.chat_stream(**cohere_params)
                 else:
                     response = client.chat(**cohere_params)
+
             except CohereRateLimitError as e:
                 raise RuntimeError(f"Cohere exception occurred: {e}")
             else:
@@ -303,6 +304,15 @@ def extract_to_cohere_tool_results(tool_call_id: str, content_output: str, all_t
     return temp_tool_results
 
 
+def is_recent_tool_call(messages: list[Dict[str, Any]], tool_call_index: int):
+    messages_length = len(messages)
+    if tool_call_index == messages_length - 1:
+        return True
+    elif messages[tool_call_index + 1].get("role", "").lower() not in ("chatbot"):
+        return True
+    return False
+
+
 def oai_messages_to_cohere_messages(
     messages: list[Dict[str, Any]], params: Dict[str, Any], cohere_params: Dict[str, Any]
 ) -> tuple[list[dict[str, Any]], str, str]:
@@ -322,7 +332,7 @@ def oai_messages_to_cohere_messages(
 
     cohere_messages = []
     preamble = ""
-
+    cohere_tool_names = set()
     # Tools
     if "tools" in params:
         cohere_tools = []
@@ -353,6 +363,7 @@ def oai_messages_to_cohere_messages(
                 "description": tool["function"]["description"],
                 "parameter_definitions": parameters,
             }
+            cohere_tool_names.add(tool["function"]["name"] or "")
 
             cohere_tools.append(cohere_tool)
 
@@ -370,42 +381,68 @@ def oai_messages_to_cohere_messages(
     # 'content' field renamed to 'message'
     # tools go into tools parameter
     # tool_results go into tool_results parameter
-    messages_length = len(messages)
     for index, message in enumerate(messages):
 
+        if not message["content"]:
+            continue
+
         if "role" in message and message["role"] == "system":
             # System message
             if preamble == "":
                 preamble = message["content"]
             else:
                 preamble = preamble + "\n" + message["content"]
-        elif "tool_calls" in message:
+
+        elif message.get("tool_calls"):
             # Suggested tool calls, build up the list before we put it into the tool_results
-            for tool_call in message["tool_calls"]:
+            message_tool_calls = []
+            for tool_call in message["tool_calls"] or []:
+                if (not tool_call.get("function", {}).get("name")) or tool_call.get("function", {}).get(
+                    "name"
+                ) not in cohere_tool_names:
+                    new_message = {
+                        "role": "CHATBOT",
+                        "message": message.get("name") + ":" + message["content"] + str(message["tool_calls"]),
+                    }
+                    cohere_messages.append(new_message)
+                    continue
+
                 tool_calls.append(tool_call)
+                message_tool_calls.append(
+                    {
+                        "name": tool_call.get("function", {}).get("name"),
+                        "parameters": json.loads(tool_call.get("function", {}).get("arguments") or "null"),
+                    }
+                )
+
+            if not message_tool_calls:
+                continue
 
             # We also add the suggested tool call as a message
             new_message = {
                 "role": "CHATBOT",
-                "message": message["content"],
-                "tool_calls": [
-                    {
-                        "name": tool_call_.get("function", {}).get("name"),
-                        "parameters": json.loads(tool_call_.get("function", {}).get("arguments") or "null"),
-                    }
-                    for tool_call_ in message["tool_calls"]
-                ],
+                "message": message.get("name") + ":" + message["content"],
+                "tool_calls": message_tool_calls,
             }
 
             cohere_messages.append(new_message)
         elif "role" in message and message["role"] == "tool":
             if not (tool_call_id := message.get("tool_call_id")):
                 continue
 
-            # Convert the tool call to a result
             content_output = message["content"]
+            if tool_call_id not in [tool_call["id"] for tool_call in tool_calls]:
+
+                new_message = {
+                    "role": "CHATBOT",
+                    "message": content_output,
+                }
+                cohere_messages.append(new_message)
+                continue
+
+            # Convert the tool call to a result
             tool_results_chat_turn = extract_to_cohere_tool_results(tool_call_id, content_output, tool_calls)
-            if (index == messages_length - 1) or (messages[index + 1].get("role", "").lower() in ("user", "tool")):
+            if is_recent_tool_call(messages, index):
                 # If the tool call is the last message or the next message is a user/tool message, this is a recent tool call.
                 # So, we pass it into tool_results.
                 tool_results.extend(tool_results_chat_turn)
@@ -420,7 +457,7 @@ def oai_messages_to_cohere_messages(
             # Standard text message
             new_message = {
                 "role": "USER" if message["role"] == "user" else "CHATBOT",
-                "message": message["content"],
+                "message": message.get("name") + ":" + message.get("content"),
             }
 
             cohere_messages.append(new_message)
@@ -436,7 +473,7 @@ def oai_messages_to_cohere_messages(
         # So, we add a CHATBOT 'continue' message, if so.
         # Changed key from "content" to "message" (jaygdesai/autogen_Jay)
         if cohere_messages[-1]["role"].lower() == "user":
-            cohere_messages.append({"role": "CHATBOT", "message": "Please continue."})
+            cohere_messages.append({"role": "CHATBOT", "message": "Please go ahead and follow the instructions!"})
 
         # We return a blank message when we have tool results
         # TODO: Check what happens if tool_results aren't the latest message
@@ -449,7 +486,7 @@ def oai_messages_to_cohere_messages(
         if cohere_messages[-1]["role"] == "USER":
             return cohere_messages[0:-1], preamble, cohere_messages[-1]["message"]
         else:
-            return cohere_messages, preamble, "Please continue."
+            return cohere_messages, preamble, "Please go ahead and follow the instructions!"
 
 
 def calculate_cohere_cost(input_tokens: int, output_tokens: int, model: str) -> float: