Merge pull request #414 from ag2ai/add-lint-rules

Add linting rules to ruff check and format code with ruff format

.pre-commit-config.yaml

@@ -23,17 +23,16 @@ repos:
     - id: trailing-whitespace
     - id: end-of-file-fixer
     - id: no-commit-to-branch
-  - repo: https://github.com/psf/black
-    rev: 24.4.2
-    hooks:
-    - id: black
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.8
+  - repo: local
     hooks:
-      - id: ruff
-        types_or: [ python, pyi, jupyter ]
-        args: ["--fix", "--ignore=E402"]
-        exclude: notebook/agentchat_databricks_dbrx.ipynb
+      - id: lint
+        name: linting and formatting
+        entry: "scripts/pre-commit-lint.sh"
+        language: python
+        #        language_version: python3.9
+        types: [python]
+        require_serial: true
+        verbose: true
   - repo: https://github.com/codespell-project/codespell
     rev: v2.3.0
     hooks:
@@ -80,7 +79,3 @@ repos:
           notebook/.* |
           website/.*
         )$
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.8.5
-    hooks:
-      - id: nbqa-black

autogen/_pydantic.py

@@ -77,11 +77,11 @@ def type2schema(t: Any) -> JsonSchemaValue:
             return {"anyOf": [type2schema(tt) for tt in get_args(t)]}
         # we need to support both syntaxes for Tuple
         elif get_origin(t) in [Tuple, tuple]:
-            prefixItems = [type2schema(tt) for tt in get_args(t)]
+            prefix_items = [type2schema(tt) for tt in get_args(t)]
             return {
-                "maxItems": len(prefixItems),
-                "minItems": len(prefixItems),
-                "prefixItems": prefixItems,
+                "maxItems": len(prefix_items),
+                "minItems": len(prefix_items),
+                "prefixItems": prefix_items,
                 "type": "array",
             }
         else:

autogen/agentchat/contrib/capabilities/generate_images.py

@@ -7,8 +7,8 @@
 import re
 from typing import Any, Dict, List, Literal, Optional, Protocol, Tuple, Union
 
-from openai import OpenAI
 from PIL.Image import Image
+from openai import OpenAI
 
 from autogen import Agent, ConversableAgent, code_utils
 from autogen.agentchat.contrib import img_utils

autogen/agentchat/contrib/capabilities/vision_capability.py

@@ -155,7 +155,7 @@ def process_last_received_message(self, content: Union[str, list[dict]]) -> str:
             ```python
             content = [
                 {"type": "text", "text": "What's weather in this cool photo:"},
-                {"type": "image_url", "image_url": {"url": "http://example.com/photo.jpg"}}
+                {"type": "image_url", "image_url": {"url": "http://example.com/photo.jpg"}},
             ]
             ```
             Output: "What's weather in this cool photo: `<img http://example.com/photo.jpg>` in case you can not see, the caption of this image is:

autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_youtube_caption.py

@@ -4,7 +4,10 @@
 # alternative api: https://rapidapi.com/omarmhaimdat/api/youtube-v2
 
 
-def get_youtube_caption(videoId):
+from typing import Any
+
+
+def get_youtube_caption(video_id: str) -> Any:
     """
     Retrieves the captions for a YouTube video.
 
@@ -21,13 +24,13 @@ def get_youtube_caption(videoId):
 
     import requests
 
-    RAPID_API_KEY = os.environ["RAPID_API_KEY"]
-    video_url = f"https://www.youtube.com/watch?v={videoId}"
+    rapid_api_key = os.environ["RAPID_API_KEY"]
+    video_url = f"https://www.youtube.com/watch?v={video_id: str}"
     url = "https://youtube-transcript3.p.rapidapi.com/api/transcript-with-url"
 
     querystring = {"url": video_url, "lang": "en", "flat_text": "true"}
 
-    headers = {"X-RapidAPI-Key": RAPID_API_KEY, "X-RapidAPI-Host": "youtube-transcript3.p.rapidapi.com"}
+    headers = {"X-RapidAPI-Key": rapid_api_key, "X-RapidAPI-Host": "youtube-transcript3.p.rapidapi.com"}
 
     response = requests.get(url, headers=headers, params=querystring)
     response = response.json()

autogen/agentchat/contrib/captainagent/tools/math/count_distinct_permutations.py

@@ -12,7 +12,7 @@ def count_distinct_permutations(sequence):
         int: The number of distinct permutations.
 
     Example:
-        >>> count_distinct_permutations('aab')
+        >>> count_distinct_permutations("aab")
         3
         >>> count_distinct_permutations([1, 2, 2])
         3

autogen/agentchat/contrib/img_utils.py

@@ -353,7 +353,7 @@ def num_tokens_from_gpt_image(
     Examples:
     --------
     >>> from PIL import Image
-    >>> img = Image.new('RGB', (2500, 2500), color = 'red')
+    >>> img = Image.new("RGB", (2500, 2500), color="red")
     >>> num_tokens_from_gpt_image(img, model="gpt-4-vision")
     765
     """

autogen/agentchat/contrib/llamaindex_conversable_agent.py

@@ -87,9 +87,9 @@ def _generate_oai_reply(
         """Generate a reply using autogen.oai."""
         user_message, history = self._extract_message_and_history(messages=messages, sender=sender)
 
-        chatResponse: AgentChatResponse = self._llama_index_agent.chat(message=user_message, chat_history=history)
+        chat_response: AgentChatResponse = self._llama_index_agent.chat(message=user_message, chat_history=history)
 
-        extracted_response = chatResponse.response
+        extracted_response = chat_response.response
 
         return (True, extracted_response)
 
@@ -102,11 +102,11 @@ async def _a_generate_oai_reply(
         """Generate a reply using autogen.oai."""
         user_message, history = self._extract_message_and_history(messages=messages, sender=sender)
 
-        chatResponse: AgentChatResponse = await self._llama_index_agent.achat(
+        chat_response: AgentChatResponse = await self._llama_index_agent.achat(
             message=user_message, chat_history=history
         )
 
-        extracted_response = chatResponse.response
+        extracted_response = chat_response.response
 
         return (True, extracted_response)
 

autogen/agentchat/contrib/math_user_proxy_agent.py

@@ -402,6 +402,7 @@ class Config:
             extra = Extra.forbid
 
     @root_validator(skip_on_failure=True)
+    @classmethod
     def validate_environment(cls, values: dict) -> dict:
         """Validate that api key and python package exists in environment."""
         wolfram_alpha_appid = get_from_dict_or_env(values, "wolfram_alpha_appid", "WOLFRAM_ALPHA_APPID")

autogen/agentchat/contrib/reasoning_agent.py

@@ -43,7 +43,6 @@
 
 
 class ThinkNode:
-
     def __init__(self, content: str, parent: Optional["ThinkNode"] = None) -> None:
         """A node in a tree structure representing a step in the reasoning process.
 
@@ -623,7 +622,8 @@ def _mtcs_reply(self, prompt, ground_truth=""):
                 # More intensive analysis is needed in the future.
                 choices_weights = [
                     # exploitation term +
-                    (child.value / (child.visits + EPSILON)) +
+                    (child.value / (child.visits + EPSILON))
+                    +
                     # exploration term
                     self._exploration_constant
                     * math.sqrt(2 * math.log(node.visits + EPSILON) / (child.visits + EPSILON))

autogen/agentchat/contrib/retrieve_user_proxy_agent.py

@@ -325,16 +325,16 @@ def _init_db(self):
         if not self._vector_db:
             return
 
-        IS_TO_CHUNK = False  # whether to chunk the raw files
+        is_to_chunk = False  # whether to chunk the raw files
         if self._new_docs:
-            IS_TO_CHUNK = True
+            is_to_chunk = True
         if not self._docs_path:
             try:
                 self._vector_db.get_collection(self._collection_name)
                 logger.warning(f"`docs_path` is not provided. Use the existing collection `{self._collection_name}`.")
                 self._overwrite = False
                 self._get_or_create = True
-                IS_TO_CHUNK = False
+                is_to_chunk = False
             except ValueError:
                 raise ValueError(
                     "`docs_path` is not provided. "
@@ -346,16 +346,16 @@ def _init_db(self):
                 self._vector_db.get_collection(self._collection_name)
                 logger.info(f"Use the existing collection `{self._collection_name}`.", color="green")
             except ValueError:
-                IS_TO_CHUNK = True
+                is_to_chunk = True
         else:
-            IS_TO_CHUNK = True
+            is_to_chunk = True
 
         self._vector_db.active_collection = self._vector_db.create_collection(
             self._collection_name, overwrite=self._overwrite, get_or_create=self._get_or_create
         )
 
         docs = None
-        if IS_TO_CHUNK:
+        if is_to_chunk:
             if self.custom_text_split_function is not None:
                 chunks, sources = split_files_to_chunks(
                     get_files_from_dir(self._docs_path, self._custom_text_types, self._recursive),

autogen/agentchat/contrib/swarm_agent.py

@@ -38,7 +38,7 @@ class AfterWorkOption(Enum):
 
 
 @dataclass
-class AFTER_WORK:
+class AFTER_WORK:  # noqa: N801
     """Handles the next step in the conversation when an agent doesn't suggest a tool call or a handoff
 
     Args:
@@ -55,7 +55,7 @@ def __post_init__(self):
 
 
 @dataclass
-class ON_CONDITION:
+class ON_CONDITION:  # noqa: N801
     """Defines a condition for transitioning to another agent or nested chats
 
     Args:
@@ -86,7 +86,7 @@ def __post_init__(self):
 
 
 @dataclass
-class UPDATE_SYSTEM_MESSAGE:
+class UPDATE_SYSTEM_MESSAGE:  # noqa: N801
     """Update the agent's system message before they reply
 
     Args:
@@ -661,7 +661,6 @@ def register_update_agent_state_before_reply(self, functions: Optional[Union[lis
 
         for func in functions:
             if isinstance(func, UPDATE_SYSTEM_MESSAGE):
-
                 # Wrapper function that allows this to be used in the update_agent_state hook
                 # Its primary purpose, however, is just to update the agent's system message
                 # Outer function to create a closure with the update function
@@ -732,7 +731,6 @@ def transfer_to_agent_name() -> SwarmAgent:
                 ), "Invalid After Work value"
                 self.after_work = transit
             elif isinstance(transit, ON_CONDITION):
-
                 if isinstance(transit.target, SwarmAgent):
                     # Transition to agent
 
@@ -809,15 +807,13 @@ def generate_swarm_tool_reply(
 
         message = messages[-1]
         if "tool_calls" in message:
-
             tool_call_count = len(message["tool_calls"])
 
             # Loop through tool calls individually (so context can be updated after each function call)
             next_agent = None
             tool_responses_inner = []
             contents = []
             for index in range(tool_call_count):
-
                 # Deep copy to ensure no changes to messages when we insert the context variables
                 message_copy = copy.deepcopy(message)
 
@@ -834,7 +830,6 @@ def generate_swarm_tool_reply(
                         # Inject the context variables into the tool call if it has the parameter
                         sig = signature(func)
                         if __CONTEXT_VARIABLES_PARAM_NAME__ in sig.parameters:
-
                             current_args = json.loads(tool_call["function"]["arguments"])
                             current_args[__CONTEXT_VARIABLES_PARAM_NAME__] = self._context_variables
                             tool_call["function"]["arguments"] = json.dumps(current_args)

autogen/agentchat/contrib/tool_retriever.py

@@ -135,7 +135,7 @@ class LocalExecutorWithTools(CodeExecutor):
     ag2_tool = interop.convert_tool(tool=langchain_tool, type="langchain")
 
     # `ag2_tool.name` is wikipedia
-    local_executor = LocalExecutorWithTools(tools=[ag2_tool], work_dir='./')
+    local_executor = LocalExecutorWithTools(tools=[ag2_tool], work_dir="./")
 
     code = '''
     result = wikipedia(tool_input={"query":"Christmas"})
@@ -161,13 +161,13 @@ def code_extractor(self) -> CodeExtractor:
         """(Experimental) Export a code extractor that can be used by an agent."""
         return MarkdownCodeExtractor()
 
-    def __init__(self, tools: Optional[List[Tool]] = None, work_dir: Union[Path, str] = Path(".")):
+    def __init__(self, tools: Optional[list[Tool]] = None, work_dir: Union[Path, str] = Path(".")):
         self.tools = tools if tools is not None else []
         self.work_dir = work_dir
         if not os.path.exists(work_dir):
             os.makedirs(work_dir, exist_ok=True)
 
-    def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CodeResult:
+    def execute_code_blocks(self, code_blocks: list[CodeBlock]) -> CodeResult:
         """Execute code blocks and return the result.
 
         Args:

autogen/agentchat/conversable_agent.py

@@ -47,8 +47,8 @@
     ConversableAgentUsageSummaryMessage,
     ConversableAgentUsageSummaryNoCostIncurredMessage,
     ExecuteCodeBlockMessage,
-    ExecutedFunctionMessage,
     ExecuteFunctionMessage,
+    ExecutedFunctionMessage,
     GenerateCodeExecutionReplyMessage,
     TerminationAndHumanReplyMessage,
     UsingAutoReplyMessage,
@@ -435,7 +435,7 @@ def _get_chats_to_run(
                 message = last_msg
             if callable(message):
                 message = message(recipient, messages, sender, config)
-            # We only run chat that has a valid message. NOTE: This is prone to change dependin on applications.
+            # We only run chat that has a valid message. NOTE: This is prone to change depending on applications.
             if message:
                 current_c["message"] = message
                 chat_to_run.append(current_c)
@@ -783,9 +783,7 @@ def send(
         ```python
         {
             "content": lambda context: context["use_tool_msg"],
-            "context": {
-                "use_tool_msg": "Use tool X if they are relevant."
-            }
+            "context": {"use_tool_msg": "Use tool X if they are relevant."},
         }
         ```
                     Next time, one agent can send a message B with a different "use_tool_msg".
@@ -833,9 +831,7 @@ async def a_send(
         ```python
         {
             "content": lambda context: context["use_tool_msg"],
-            "context": {
-                "use_tool_msg": "Use tool X if they are relevant."
-            }
+            "context": {"use_tool_msg": "Use tool X if they are relevant."},
         }
         ```
                     Next time, one agent can send a message B with a different "use_tool_msg".
@@ -2721,7 +2717,6 @@ def _decorator(func_or_tool: Union[F, Tool]) -> Tool:
         return _decorator
 
     def _register_for_llm(self, tool: Tool, api_style: Literal["tool", "function"]) -> None:
-
         # register the function to the agent if there is LLM config, raise an exception otherwise
         if self.llm_config is None:
             raise RuntimeError("LLM config must be setup before registering a function for LLM.")

autogen/agentchat/groupchat.py

@@ -1107,9 +1107,7 @@ def last_speaker(self) -> Agent:
 
         def print_messages(recipient, messages, sender, config):
             # Print the message immediately
-            print(
-                f"Sender: {sender.name} | Recipient: {recipient.name} | Message: {messages[-1].get('content')}"
-            )
+            print(f"Sender: {sender.name} | Recipient: {recipient.name} | Message: {messages[-1].get('content')}")
             print(f"Real Sender: {sender.last_speaker.name}")
             assert sender.last_speaker.name in messages[-1].get("content")
             return False, None  # Required to ensure the agent communication flow continues
@@ -1119,9 +1117,7 @@ def print_messages(recipient, messages, sender, config):
         agent_b = ConversableAgent("agent B", default_auto_reply="I'm agent B.")
         agent_c = ConversableAgent("agent C", default_auto_reply="I'm agent C.")
         for agent in [agent_a, agent_b, agent_c]:
-            agent.register_reply(
-                [ConversableAgent, None], reply_func=print_messages, config=None
-            )
+            agent.register_reply([ConversableAgent, None], reply_func=print_messages, config=None)
         group_chat = GroupChat(
             [agent_a, agent_b, agent_c],
             messages=[],
@@ -1130,9 +1126,7 @@ def print_messages(recipient, messages, sender, config):
             allow_repeat_speaker=True,
         )
         chat_manager = GroupChatManager(group_chat)
-        groupchat_result = agent_a.initiate_chat(
-            chat_manager, message="Hi, there, I'm agent A."
-        )
+        groupchat_result = agent_a.initiate_chat(chat_manager, message="Hi, there, I'm agent A.")
         ```
         """
         return self._last_speaker