ag2 0.9.1a1__py3-none-any.whl → 0.9.2__py3-none-any.whl

autogen/mcp/mcp_proxy/operation_grouping.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+from pathlib import Path
+from typing import Dict, List
+
+from autogen.import_utils import optional_import_block
+
+with optional_import_block() as result:
+    from fastapi_code_generator.parser import OpenAPIParser, Operation
+    from fastapi_code_generator.visitor import Visitor
+
+from pydantic import BaseModel
+
+from autogen.agentchat.conversable_agent import ConversableAgent
+from autogen.llm_config import LLMConfig
+
+
+class Group(BaseModel):
+    name: str
+    description: str
+
+
+class GroupSuggestions(BaseModel):
+    groups: list[Group]
+
+
+class GroupNames(BaseModel):
+    groups: list[str]
+
+
+logger = logging.getLogger(__name__)
+
+GROUP_DISCOVERY_MESSAGE = (
+    "You are a senior Python engineer. You will be shown a batch of API functions. "
+    "These functions are not guaranteed to be related — your task is to analyze them individually and find meaningful groups *within the batch*.\n\n"
+    "You should propose a list of group names, and for each group, provide a short description of the kind of functions it contains.\n\n"
+    "How to group:\n"
+    "- Focus on what the function operates on — e.g., functions that manipulate a board go in 'board_operations', functions related to users in 'user_management', etc.\n"
+    "- Do NOT assume the entire batch forms a single group.\n"
+    "- Do NOT use generic categories like 'misc', 'utils', or 'general'.\n"
+    "- Favor *granular but meaningful* groups. For instance, 'user_profile_handling' and 'user_authentication' could be separate if their purposes are distinct.\n\n"
+    "Formatting:\n"
+    "- Group names must be short, descriptive, and in snake_case.\n"
+    "- Return a list of group names and a 1-2 sentence description of what each group includes.\n"
+)
+
+
+GROUP_ASSIGNMENT_MESSAGE = (
+    "You are a senior Python engineer. You will be given a function description "
+    "and a list of possible groups. Choose the most suitable groups for the function.\n"
+    "- A function can belong to multiple groups.\n"
+    "- The groups should be relevant to the function's purpose and functionality.\n"
+    "- You MUST choose a group from the possible groups, you cannot respond with empty grouping any equivalent.\n"
+)
+
+
+def chunk_list(items: List, size: int) -> List[List]:
+    return [items[i : i + size] for i in range(0, len(items), size)]
+
+
+def discover_groups(operations: List["Operation"], chunk_size: int = 30) -> Dict[str, str]:
+    llm_config = LLMConfig.get_current_llm_config().copy()
+
+    for config in llm_config.config_list:
+        config.response_format = GroupSuggestions
+
+    with llm_config:
+        agent = ConversableAgent(name="group_discovery_agent", system_message=GROUP_DISCOVERY_MESSAGE)
+        groups = {}
+
+        for i, chunk in enumerate(chunk_list(operations, chunk_size)):
+            func_descriptions = [f"- {op.function_name}: {op.summary} (args: {op.arguments})" for op in chunk]
+            message = "Here are some functions:\n" + "\n".join(func_descriptions)
+
+            response = agent.run(message=message, max_turns=1, user_input=False)
+
+            for event in response.events:
+                if event.type == "text" and event.content.sender == "group_discovery_agent":
+                    # Naively parse "group_name: description" from text block
+                    new_groups = GroupSuggestions.model_validate_json(event.content.content).groups
+                    groups.update(new_groups)
+
+    logger.warning("Discovered groups: %s", groups)
+
+    # Remove duplicates
+    with llm_config:
+        agent = ConversableAgent(name="group_refining_agent", system_message=GROUP_DISCOVERY_MESSAGE)
+
+        message = (
+            "You need to refine the group names and descriptions to ensure they are unique.\n"
+            "Here are the groups:\n" + "\n".join([f"- {name}: {desc}" for name, desc in groups.items()])
+        )
+        response = agent.run(message=message, max_turns=1, user_input=False)
+        for event in response.events:
+            if event.type == "text" and event.content.sender == "group_refining_agent":
+                # Naively parse "group_name: description" from text block
+                refined_groups = json.loads(event.content.content)
+
+    return refined_groups
+
+
+def assign_operation_to_group(operation: "Operation", groups: dict[str, str]) -> str:
+    llm_config = LLMConfig.get_current_llm_config().copy()
+
+    for config in llm_config.config_list:
+        config.response_format = GroupNames
+
+    with llm_config:
+        agent = ConversableAgent(name="group_assignment_agent", system_message=GROUP_ASSIGNMENT_MESSAGE)
+
+        message = (
+            "Function summary:\n"
+            f"{operation.summary}\n\n"
+            f"Arguments: {operation.arguments}\n\n"
+            f"Available groups: {json.dumps(groups)}\n\n"
+            "What group should this function go in?"
+        )
+
+        response = agent.run(message=message, max_turns=1, user_input=True)
+
+        groups = []
+        for event in response.events:
+            if event.type == "text" and event.content.sender == "group_assignment_agent":
+                groups = GroupNames.model_validate_json(event.content.content).groups
+
+        return groups
+
+
+def refine_group_names(groups: Dict[str, str]) -> Dict[str, str]:
+    # Optional: normalize names, merge similar ones (e.g., using embeddings or string similarity)
+    # Placeholder for now:
+    return groups
+
+
+def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> Dict[str, object]:
+    operations = sorted(parser.operations.values(), key=lambda op: op.path)
+
+    # ---- PASS 1: DISCOVER GROUPS ----
+    logger.warning("Discovering groups...")
+    discovered_groups = discover_groups(operations)
+    logger.warning("Discovered groups: %s", discovered_groups)
+
+    # ---- PASS 2: ASSIGN OPERATIONS TO GROUPS ----
+    logger.warning("Assigning operations to groups...")
+    for op in operations:
+        logger.warning("Assigning operation %s to groups...", op.function_name)
+        groups = assign_operation_to_group(op, discovered_groups)
+        op.tags = groups
+        logger.warning("Assigned groups: %s", groups)
+
+    return {"operations": operations}
+
+
+visit: "Visitor" = custom_visitor

autogen/mcp/mcp_proxy/operation_renaming.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from pathlib import Path
+from typing import Dict, List
+
+from autogen.import_utils import optional_import_block
+
+with optional_import_block() as result:
+    from fastapi_code_generator.parser import OpenAPIParser, Operation
+    from fastapi_code_generator.visitor import Visitor
+
+from autogen.agentchat.conversable_agent import ConversableAgent
+
+logger = logging.getLogger(__name__)
+
+# System prompt to guide the AI agent in naming functions
+SYSTEM_MESSAGE = (
+    "You are a helpful expert Python programmer. Your task is to generate a clear, concise, "
+    "and descriptive name for a Python function based on a user-provided summary.\n"
+    "- Only return the new function name.\n"
+    "- The name must be fewer than 64 characters.\n"
+    "- It should reflect the purpose of the function as described.\n"
+    "- You will be provided with a list of already-taken names, which must not be reused.\n"
+    "- The function name should be in snake_case.\n"
+)
+
+
+def validate_function_name(name: str, taken_names: List[str]) -> str:
+    """
+    Validate the generated function name against length, format, and uniqueness constraints.
+
+    Returns:
+        'exit' if the name is valid, or an error message string otherwise.
+    """
+    if len(name) > 64:
+        return "Function name is too long. Please provide a shorter name."
+    if not name.islower() or " " in name:
+        return "Function name must be in snake_case."
+    if name in taken_names:
+        return f"Function name is already taken. Please provide a different name. Taken names: {', '.join(taken_names)}"
+    return "exit"
+
+
+def get_new_function_name(operation: "Operation", taken_names: List[str]) -> str:
+    """
+    Ask an AI agent to generate a new function name for a given OpenAPI operation.
+
+    Args:
+        operation: The OpenAPI operation that needs renaming.
+        taken_names: A list of names already used, to avoid collisions.
+
+    Returns:
+        A new, validated function name.
+    """
+    agent = ConversableAgent(
+        name="helpful_agent",
+        system_message=SYSTEM_MESSAGE,
+    )
+
+    response = agent.run(
+        message=(
+            "How would you name this function? \n"
+            f"Info:\n"
+            f"- old function name: {operation.function_name}\n"
+            f"- function summary: {operation.summary}\n"
+            f"- function arguments: {operation.arguments}\n"
+        ),
+        user_input=True,
+    )
+
+    proposed_name = None
+
+    for event in response.events:
+        if event.type == "text" and event.content.sender == "helpful_agent":
+            proposed_name = event.content.content.strip()
+        elif event.type == "input_request":
+            reply = (
+                validate_function_name(proposed_name, taken_names)
+                if proposed_name
+                else "Please provide a function name."
+            )
+            event.content.respond(reply)
+
+    logger.warning(f"Renamed operation '{operation.function_name}' to '{response.summary}'.")
+    return response.summary
+
+
+def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> Dict[str, object]:
+    """
+    Visits and optionally renames operations in the OpenAPI parser.
+
+    Args:
+        parser: An OpenAPIParser instance containing API operations.
+        model_path: Path to the model (not used in this implementation).
+
+    Returns:
+        A dictionary containing the updated list of operations.
+    """
+    operations = sorted(parser.operations.values(), key=lambda op: op.path)
+    taken_names = []
+
+    for op in operations:
+        if len(op.function_name) > 64:
+            new_name = get_new_function_name(op, taken_names)
+            op.function_name = new_name
+        taken_names.append(op.function_name)
+
+    return {"operations": operations}
+
+
+visit: "Visitor" = custom_visitor

autogen/mcp/mcp_proxy/patch_fastapi_code_generator.py

@@ -0,0 +1,98 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+import logging
+import re
+from functools import cached_property, wraps
+from typing import Any
+
+from ...import_utils import optional_import_block, require_optional_import
+
+with optional_import_block():
+    import yaml
+
+from autogen.import_utils import optional_import_block
+
+with optional_import_block() as result:
+    from fastapi_code_generator import __main__ as fastapi_code_generator_main
+    from fastapi_code_generator.parser import OpenAPIParser, Operation
+
+SUCCESFUL_IMPORT = result.is_successful
+
+logger = logging.getLogger(__name__)
+
+
+def patch_parse_schema() -> None:
+    org_parse_schema = OpenAPIParser.parse_schema
+
+    @wraps(org_parse_schema)
+    def my_parse_schema(*args: Any, **kwargs: Any) -> Any:
+        data_type = org_parse_schema(*args, **kwargs)
+        if data_type.reference and data_type.reference.duplicate_name:
+            data_type.reference.name = data_type.reference.duplicate_name
+        return data_type
+
+    OpenAPIParser.parse_schema = my_parse_schema
+    logger.info("Patched OpenAPIParser.parse_schema")
+
+
+def snakecase(string: str) -> str:
+    string = re.sub(r"[\-\.\s]", "_", str(string))
+    if not string:
+        return string
+    return string[0].lower() + re.sub(r"[A-Z]", lambda matched: "_" + (matched.group(0)), string[1:]).lower()
+
+
+def patch_function_name_parsing() -> None:
+    def function_name(self: Operation) -> str:
+        if self.operationId:
+            name: str = re.sub(r"[/{=}]", "_", self.operationId)
+        else:
+            path = re.sub(r"[/{=]", "_", self.snake_case_path).replace("}", "")
+            name = f"{self.type}{path}"
+
+        return snakecase(name)  # type: ignore[no-any-return]
+
+    Operation.function_name = cached_property(function_name)
+    Operation.function_name.__set_name__(Operation, "function_name")
+
+    logger.info("Patched Operation.function_name")
+
+
+@require_optional_import(["yaml"], "mcp-proxy-gen")
+def patch_generate_code() -> None:
+    # Save reference to the original generate_code function
+    org_generate_code = fastapi_code_generator_main.generate_code
+
+    @wraps(org_generate_code)
+    def patched_generate_code(*args: Any, **kwargs: Any) -> Any:
+        try:
+            input_text: str = kwargs["input_text"]
+
+            json_spec = yaml.safe_load(input_text)
+
+            schemas_with_dots = sorted(
+                [name for name in json_spec.get("components", {}).get("schemas", {}) if "." in name],
+                key=len,
+                reverse=True,  # Sort by length in descending order
+            )
+
+            for schema_name in schemas_with_dots:
+                new_schema_name = schema_name.replace(".", "_")
+                input_text = input_text.replace(schema_name, new_schema_name)
+
+            kwargs["input_text"] = input_text
+
+        except Exception as e:
+            print(
+                f"Patched fastapi_code_generator.__main__.generate_code raised: {e}, passing untouched arguments to original generate_code"
+            )
+            logger.info(
+                f"Patched fastapi_code_generator.__main__.generate_code raised: {e}, passing untouched arguments to original generate_code"
+            )
+
+        return org_generate_code(*args, **kwargs)
+
+    fastapi_code_generator_main.generate_code = patched_generate_code
+
+    logger.info("Patched fastapi_code_generator.__main__.generate_code")