feat(prompts_sample): fixed minor issues from reviews

Author: giuliastf

samples/mcp-refactoring-assistant/.agent/SDK_REFERENCE.md

@@ -398,7 +398,7 @@ Guardrails service
 
 ```python
 # Validate input text using the provided guardrail.
-sdk.guardrails.evaluate_guardrail(input_data: str | dict[str, Any], guardrail: uipath.platform.guardrails.guardrails.BuiltInValidatorGuardrail) -> uipath.core.guardrails.guardrails.GuardrailValidationResult
+sdk.guardrails.evaluate_guardrail(input_data: str | dict[str, Any], guardrail: Annotated[Union[uipath.platform.guardrails.guardrails.CustomGuardrail, uipath.platform.guardrails.guardrails.BuiltInValidatorGuardrail], FieldInfo(annotation=NoneType, required=True, discriminator='guardrail_type')]) -> uipath.platform.guardrails.guardrails.GuardrailValidationResult
 
 ```
 

samples/mcp-refactoring-assistant/entry-points.json

@@ -4,7 +4,7 @@
     "entryPoints": [
         {
             "filePath": "agent",
-            "uniqueId": "426f3cb2-ac10-4969-840e-399932250e3a",
+            "uniqueId": "d016126e-57b7-4493-93de-d834c86b9b70",
             "type": "agent",
             "input": {
                 "type": "object",

samples/mcp-refactoring-assistant/graph.py

@@ -4,48 +4,55 @@
 import json
 import sys
 from pathlib import Path
-from typing import Any, Optional
 
 from langchain.agents import create_agent
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langgraph.graph import END, START, MessagesState, StateGraph
 from pydantic import BaseModel
+
 from uipath_langchain.chat import UiPathChat
 
 model = UiPathChat(model="gpt-4o-2024-08-06", temperature=0.7)
 
 server_path = Path(__file__).parent / "server.py"
 
-_client = None
-_react_agent = None
+manager = None
+
+
+class MCPManager:
+    def __init__(self):
+        self._client: MultiServerMCPClient | None = None
+        self._react_agent = None
+
+    async def get_client(self):
+        if self._client is None:
+            self._client = MultiServerMCPClient({
+                "code-refactoring": {
+                    "command": sys.executable,
+                    "args": [str(server_path)],
+                    "transport": "stdio",
+                },
+            })
+        return self._client
 
+    async def get_agent(self):
+        if self._react_agent is None:
+            client = await self.get_client()
+            tools = await client.get_tools()
 
-async def _get_client():
-    """Get or initialize the MCP client."""
-    global _client
-    if _client is None:
-        _client = MultiServerMCPClient({
-            "code-refactoring": {
-                "command": sys.executable,
-                "args": [str(server_path)],
-                "transport": "stdio",
-            },
-        })
-    return _client
+            model_with_tools = model.bind_tools(
+                tools,
+                tool_choice={"type": "function", "function": {"name": "get_refactoring_guide"}}
+            )
+            self._react_agent = create_agent(model_with_tools, tools=tools)
+        return self._react_agent
 
 
-async def _get_agent():
-    """Get or initialize the ReAct agent with MCP tools."""
-    global _react_agent
-    if _react_agent is None:
-        client = await _get_client()
-        tools = await client.get_tools()
-        _react_agent = create_agent(model, tools=tools)
-    return _react_agent
+manager = MCPManager()
 
 
-def _try_parse_json(value: Any) -> Optional[dict]:
+def _try_parse_json(value) -> dict | None:
     """Robustly parse JSON from various formats (dict, str, list)."""
     if value is None:
         return None
@@ -95,27 +102,25 @@ class State(BaseModel):
 
 async def agent_node(state: State) -> State:
     """Agent analyzes code and determines which prompt to use."""
-    react_agent = await _get_agent()
-
-    initial_msg = HumanMessage(
-        content=f"""You are a refactoring assistant.
-
-1) Analyze this code using analyze_code_complexity
-2) Detect issues using detect_code_smells
-3) Call get_refactoring_guide with:
-   - issue_type: the main issue detected
-   - code: the code to refactor
-   - complexity_info: results from step 1
-   - smells_info: results from step 2
-
-The get_refactoring_guide tool will return {{\"prompt_name\": \"...\", \"arguments\": {{...}}}} ready for the next step.
-
-CODE:
-{state.code}
-"""
+    react_agent = await manager.get_agent()
+
+    system_msg = SystemMessage(
+        content=(
+            "You are a refactoring assistant.\n\n"
+            "1) Analyze this code using analyze_code_complexity\n"
+            "2) Detect issues using detect_code_smells\n"
+            "3) Call get_refactoring_guide with:\n"
+            "   - issue_type: the main issue detected\n"
+            "   - code: the code to refactor\n"
+            "   - complexity_info: results from step 1\n"
+            "   - smells_info: results from step 2\n\n"
+            'The get_refactoring_guide tool will return {"prompt_name": "...", "arguments": {...}} ready for the next step.'
+        )
     )
 
-    messages_state = MessagesState(messages=[initial_msg])
+    user_msg = HumanMessage(content=state.code)
+
+    messages_state = MessagesState(messages=[system_msg, user_msg])
     result = await react_agent.ainvoke(messages_state)
 
     prompt_name = ""
@@ -132,19 +137,11 @@ async def agent_node(state: State) -> State:
                     break
 
     if not prompt_name:
-        return State(
-            code=state.code,
-            prompt_name="",
-            prompt_arguments={},
-        )
+        return State(code=state.code, prompt_name="", prompt_arguments={})
 
     prompt_args.setdefault('code', state.code)
 
-    return State(
-        code=state.code,
-        prompt_name=prompt_name,
-        prompt_arguments=prompt_args,
-    )
+    return State(code=state.code, prompt_name=prompt_name, prompt_arguments=prompt_args)
 
 
 async def prompt_node(state: State) -> Output:
@@ -155,7 +152,7 @@ async def prompt_node(state: State) -> Output:
                    "Please ensure the agent analyzed the code and called get_refactoring_guide."
         )
 
-    client = await _get_client()
+    client = await manager.get_client()
 
     prompt_messages = await client.get_prompt(
         "code-refactoring",
@@ -179,3 +176,4 @@ async def prompt_node(state: State) -> Output:
 
 graph = builder.compile()
 
+

samples/mcp-refactoring-assistant/server.py

@@ -1,5 +1,3 @@
-from typing import Any, Dict, List
-
 from mcp.server.fastmcp import FastMCP
 from mcp.server.fastmcp.prompts.base import Message
 from prompts import (extract_method, improve_naming, remove_duplication,
@@ -11,7 +9,7 @@
 
 # Register Code Analysis Tools
 @mcp.tool()
-def analyze_code_complexity(code: str) -> Dict[str, Any]:
+def analyze_code_complexity(code: str) -> dict:
     """Analyze Python code complexity using heuristics.
 
     Args:
@@ -24,7 +22,7 @@ def analyze_code_complexity(code: str) -> Dict[str, Any]:
 
 
 @mcp.tool()
-def detect_code_smells(code: str) -> Dict[str, Any]:
+def detect_code_smells(code: str) -> dict:
     """Detect common code smells in Python code.
 
     Args:
@@ -40,9 +38,9 @@ def detect_code_smells(code: str) -> Dict[str, Any]:
 def get_refactoring_guide(
     issue_type: str,
     code: str = "",
-    complexity_info: Dict[str, Any] = None,
-    smells_info: Dict[str, Any] = None
-) -> Dict[str, Any]:
+    complexity_info: dict | None = None,
+    smells_info: dict | None = None
+) -> dict:
     """Get guidance on which refactoring prompt to use for a specific issue.
 
     Args:
@@ -65,7 +63,7 @@ def get_refactoring_guide(
 
 # Register Refactoring Prompts
 @mcp.prompt()
-def extract_method_prompt(code: str, target_lines: str = "auto") -> List[Message]:
+def extract_method_prompt(code: str, target_lines: str = "auto") -> list[Message]:
     """Guide for extracting a method from complex code.
 
     Args:
@@ -79,7 +77,7 @@ def extract_method_prompt(code: str, target_lines: str = "auto") -> List[Message
 
 
 @mcp.prompt()
-def simplify_conditional_prompt(code: str, pattern: str = "guard_clause") -> List[Message]:
+def simplify_conditional_prompt(code: str, pattern: str = "guard_clause") -> list[Message]:
     """Guide for simplifying complex conditional logic.
 
     Args:
@@ -93,7 +91,7 @@ def simplify_conditional_prompt(code: str, pattern: str = "guard_clause") -> Lis
 
 
 @mcp.prompt()
-def remove_duplication_prompt(code: str, duplicate_blocks: str) -> List[Message]:
+def remove_duplication_prompt(code: str, duplicate_blocks: str) -> list[Message]:
     """Guide for removing code duplication.
 
     Args:
@@ -107,7 +105,7 @@ def remove_duplication_prompt(code: str, duplicate_blocks: str) -> List[Message]
 
 
 @mcp.prompt()
-def improve_naming_prompt(code: str, symbols: str, context: str = "") -> List[Message]:
+def improve_naming_prompt(code: str, symbols: str, context: str = "") -> list[Message]:
     """Guide for improving variable and function names.
 
     Args:

samples/mcp-refactoring-assistant/tools/code_analysis.py

@@ -1,10 +1,9 @@
 """Code analysis tools for detecting complexity and code smells."""
 
 import re
-from typing import Any, Dict
 
 
-def analyze_code_complexity(code: str) -> Dict[str, Any]:
+def analyze_code_complexity(code: str) -> dict:
     """Analyze Python code complexity using heuristics.
 
     Args:
@@ -64,7 +63,7 @@ def analyze_code_complexity(code: str) -> Dict[str, Any]:
     }
 
 
-def detect_code_smells(code: str) -> Dict[str, Any]:
+def detect_code_smells(code: str) -> dict:
     """Detect common code smells in Python code.
 
     Args:

samples/mcp-refactoring-assistant/tools/refactoring_guide.py

@@ -1,14 +1,12 @@
-"""Refactoring guide tool that maps code issues to refactoring prompts."""
-
-from typing import Any, Dict, Optional
+"""Refactoring guide tool that maps code issues to refactoring prompts."""
 
 
 def get_refactoring_guide(
     issue_type: str,
     code: str = "",
-    complexity_info: Optional[Dict[str, Any]] = None,
-    smells_info: Optional[Dict[str, Any]] = None
-) -> Dict[str, Any]:
+    complexity_info: dict | None = None,
+    smells_info: dict | None = None
+) -> dict:
     """Get guidance on which refactoring prompt to use for a specific issue.
 
 
@@ -101,47 +99,68 @@ def get_refactoring_guide(
 
     # Construct arguments dict based on prompt requirements
     if prompt_name == "extract_method_prompt":
-        arguments = {
-            "code": code,
-            "target_lines": "auto"  # Agent can override if needed
+        return {
+            "prompt_name": prompt_name,
+            "arguments": {
+                "code": code,
+                "target_lines": "auto"
+            },
+            "description": guide["description"],
+            "priority": guide["priority"]
         }
-    elif prompt_name == "simplify_conditional_prompt":
+
+    if prompt_name == "simplify_conditional_prompt":
         # Determine pattern based on issue type
         pattern = "guard_clause" if issue_type == "deep_nesting" else "early_return"
-        arguments = {
-            "code": code,
-            "pattern": pattern
+        return {
+            "prompt_name": prompt_name,
+            "arguments": {
+                "code": code,
+                "pattern": pattern
+            },
+            "description": guide["description"],
+            "priority": guide["priority"]
         }
-    elif prompt_name == "remove_duplication_prompt":
+
+    if prompt_name == "remove_duplication_prompt":
         # Extract duplicate info from smells_info if available
         duplicate_blocks = "auto-detected duplicates"
         if smells_info and "duplicate_code" in smells_info.get("smells", []):
             duplicate_blocks = "See code analysis for duplicate blocks"
-        arguments = {
-            "code": code,
-            "duplicate_blocks": duplicate_blocks
+        return {
+            "prompt_name": prompt_name,
+            "arguments": {
+                "code": code,
+                "duplicate_blocks": duplicate_blocks
+            },
+            "description": guide["description"],
+            "priority": guide["priority"]
         }
-    elif prompt_name == "improve_naming_prompt":
+
+    if prompt_name == "improve_naming_prompt":
         # Extract poor naming symbols from smells_info
         symbols = "x,y,z"  # Default, agent can override
         if smells_info:
             poor_names = [s.get("symbol", "") for s in smells_info.get("smells", [])
                          if s.get("type") == "poor_naming"]
             if poor_names:
                 symbols = ",".join(poor_names)
-        arguments = {
-            "code": code,
-            "symbols": symbols,
-            "context": ""  # Optional, agent can provide
+        return {
+            "prompt_name": prompt_name,
+            "arguments": {
+                "code": code,
+                "symbols": symbols,
+                "context": ""  # Optional, agent can provide
+            },
+            "description": guide["description"],
+            "priority": guide["priority"]
         }
-    else:
-        # Fallback: just provide code
-        arguments = {"code": code}
 
-    # Return structured response ready for client.get_prompt()
+    # Fallback: just provide code
     return {
         "prompt_name": prompt_name,
-        "arguments": arguments,
+        "arguments": {"code": code},
         "description": guide["description"],
         "priority": guide["priority"]
     }
+