diff --git a/src/agents/common/constants.py b/src/agents/common/constants.py index 22f55ab3..eceb889a 100644 --- a/src/agents/common/constants.py +++ b/src/agents/common/constants.py @@ -44,6 +44,12 @@ KYMA_AGENT = "KymaAgent" +RESPONSE_CONVERTER = "ResponseConverter" + +NEW_YAML = "New" + +UPDATE_YAML = "Update" + SUCCESS_CODE = 200 ERROR_RATE_LIMIT_CODE = 429 diff --git a/src/agents/common/response_converter.py b/src/agents/common/response_converter.py new file mode 100644 index 00000000..e04c4918 --- /dev/null +++ b/src/agents/common/response_converter.py @@ -0,0 +1,272 @@ +import re +from typing import Any, Protocol + +import yaml +from langchain_core.messages import AIMessage +from langgraph.constants import END + +from agents.common.constants import ( + FINALIZER, + MESSAGES, + NEW_YAML, + NEXT, + UPDATE_YAML, +) +from utils.logging import get_logger + +logger = get_logger(__name__) + + +class IResponseConverter(Protocol): + """Protocol for IResponseConverter.""" + + def convert_final_response(self, state: dict[str, Any]) -> dict[str, Any]: + """ + Main conversion method that orchestrates the entire YAML to HTML conversion process. + """ + ... + + +class ResponseConverter: + """ + A class that handles the conversion of YAML responses into HTML format with resource links. + This converter processes both new and update YAML configurations and generates appropriate + resource links based on the YAML content. + """ + + def __init__(self): + # Regular expression patterns to extract YAML blocks + self.new_yaml_pattern = r"\n(.*?)\n" + self.update_yaml_pattern = r"\n(.*?)\n" + + def _extract_yaml(self, finalizer_response: str) -> tuple[list[str], list[str]]: + """ + Extract YAML code blocks from finalizer response using regex patterns. + + Args: + finalizer_response: Response of the finalizer node + + Returns: + Tuple containing lists of new and update YAML blocks + """ + + # Find all YAML blocks marked for new resources + new_yaml_blocks = re.findall( + self.new_yaml_pattern, finalizer_response, re.DOTALL + ) + + # Find all YAML blocks marked for updating existing resources + update_yaml_blocks = re.findall( + self.update_yaml_pattern, finalizer_response, re.DOTALL + ) + + return new_yaml_blocks, update_yaml_blocks + + def _parse_yamls(self, yaml_config: str) -> Any | None: + """ + Parse YAML string into Python object with error handling. + Attempts two parsing methods: + 1. First check: if yaml markers available + - tries parsing after removing yaml markers + 2. Else, tries parsing the raw string + + Args: + yaml_config: YAML configuration string + + Returns: + Parsed YAML object or None if parsing fails + """ + try: + # First check: if yaml markers available + if yaml_config[:7] == "```yaml": + # parsing after removing yaml markers + parsed_yaml = yaml.safe_load(yaml_config[8:-4]) + else: + # Parse raw string + parsed_yaml = yaml.safe_load(yaml_config) + + except Exception as e: + logger.error( + f"Error while parsing the yaml : {yaml_config} , Exception : {e}" + ) + + return None + + return parsed_yaml + + def _generate_resource_link( + self, yaml_config: dict[str, Any], link_type: str + ) -> str | None: + """ + Generate resource link based on YAML metadata and link type. + + Args: + yaml_config: Parsed YAML configuration + link_type: Type of link to generate ('NEW' or 'UPDATE') + + Returns: + Generated resource link or None if required metadata is missing + """ + # Extract required metadata for link generation + try: + namespace = yaml_config["metadata"]["namespace"] + deployment_name = yaml_config["metadata"]["name"] + resource_type = yaml_config["kind"] + except Exception as e: + logger.error( + f"Error in generating link, skipping the yaml: {yaml_config} , Exception : {e}" + ) + return None + + # Generate appropriate link based on type + if link_type == NEW_YAML: + # New resource link format + return f"/namespaces/{namespace}/{resource_type}" + else: + # Update resource link format includes deployment name + return f"/namespaces/{namespace}/{resource_type}/{deployment_name}" + + def _create_html_nested_yaml( + self, yaml_config: str, resource_link: str, link_type: str + ) -> str: + """ + Create HTML structure containing YAML content and resource link. + + Args: + yaml_config: YAML configuration string + resource_link: Generated resource link + link_type: Type of link ('New' or 'Update') + + Returns: + Formatted HTML string containing YAML and link + """ + html_content = f""" +
+ {yaml_config} +
+
+ [Apply]({resource_link}) +
+ + """ + + return html_content + + def _replace_yaml_with_html( + self, + finalizer_response: str, + replacement_html_list: list[str], + yaml_type: str, + ) -> str: + """ + Replace YAML blocks in the response with corresponding HTML blocks. + + Args: + finalizer_response: Original response containing YAML blocks + replacement_html_list: List of HTML blocks to replace YAML + yaml_type: Type of YAML blocks to replace ('NEW' or 'UPDATE') + + Returns: + Modified response with YAML blocks replaced by HTML + """ + + def replace_func(match: Any) -> Any: + # Replace each match with the next HTML block from the list + if replacement_html_list: + html_content = replacement_html_list.pop(0) + return html_content + return match.group(0) + + # Select appropriate pattern based on YAML type + yaml_pattern = ( + self.new_yaml_pattern if yaml_type == NEW_YAML else self.update_yaml_pattern + ) + + # Perform the replacement + converted_response = re.sub( + yaml_pattern, replace_func, finalizer_response, flags=re.DOTALL + ) + return converted_response + + def _create_replacement_list( + self, yaml_list: list[str], yaml_type: str + ) -> list[str]: + """ + Process list of YAML configs and create corresponding HTML replacements. + + Args: + yaml_list: List of YAML configurations to process + yaml_type: Type of YAML blocks ('NEW' or 'UPDATE') + + Returns: + List of HTML replacements for YAML blocks + """ + replacement_list = [] + + for yaml_config_string in yaml_list: + # Parse YAML and generate replacement + parsed_yaml = self._parse_yamls(yaml_config_string) + if not parsed_yaml: + replacement_list.append(yaml_config_string) + continue + + # Generate resource link + generated_link = self._generate_resource_link(parsed_yaml, yaml_type) + + if generated_link: + # Create HTML if link generation successful + html_string = self._create_html_nested_yaml( + yaml_config_string, generated_link, yaml_type + ) + replacement_list.append(html_string) + else: + # Keep original if link generation fails + replacement_list.append(yaml_config_string) + + return replacement_list + + def convert_final_response(self, state: dict[str, Any]) -> dict[str, Any]: + """ + Main conversion method that orchestrates the entire YAML to HTML conversion process. + + Args: + state: Current supervisor state + + Returns: + Dictionary containing converted messages and next state + """ + finalizer_response = str(state["messages"][-1].content) + try: + # Extract all YAML blocks + new_yaml_list, update_yaml_list = self._extract_yaml(finalizer_response) + + if new_yaml_list or update_yaml_list: + # Process new resource YAML configs + replacement_list = self._create_replacement_list( + new_yaml_list, NEW_YAML + ) + finalizer_response = self._replace_yaml_with_html( + finalizer_response, replacement_list, NEW_YAML + ) + + # Process update resource YAML configs + replacement_list = self._create_replacement_list( + update_yaml_list, UPDATE_YAML + ) + finalizer_response = self._replace_yaml_with_html( + finalizer_response, replacement_list, UPDATE_YAML + ) + + except Exception as e: + logger.error(f"Error in converting final response: {e}") + + return { + MESSAGES: [ + AIMessage( + content=finalizer_response, + name=FINALIZER, + ) + ], + NEXT: END, + } diff --git a/src/agents/supervisor/agent.py b/src/agents/supervisor/agent.py index 40178b96..155a0147 100644 --- a/src/agents/supervisor/agent.py +++ b/src/agents/supervisor/agent.py @@ -19,6 +19,7 @@ NEXT, PLANNER, ) +from agents.common.response_converter import IResponseConverter, ResponseConverter from agents.common.state import Plan from agents.common.utils import create_node_output, filter_messages from agents.supervisor.prompts import ( @@ -81,12 +82,20 @@ class SupervisorAgent: members: list[str] = [] plan_parser = PydanticOutputParser(pydantic_object=Plan) - def __init__(self, models: dict[str, IModel | Embeddings], members: list[str]): + def __init__( + self, + models: dict[str, IModel | Embeddings], + members: list[str], + response_converter: IResponseConverter | None = None, + ) -> None: gpt_4o = cast(IModel, models[ModelType.GPT4O]) self.model = gpt_4o self.members = members self.parser = self._route_create_parser() + self.response_converter: IResponseConverter = ( + response_converter or ResponseConverter() + ) self._planner_chain = self._create_planner_chain(gpt_4o) self._graph = self._build_graph() @@ -235,6 +244,15 @@ async def _generate_final_response(self, state: SupervisorState) -> dict[str, An ] } + async def _get_converted_final_response( + self, state: SupervisorState + ) -> dict[str, Any]: + """Convert the generated final response.""" + + final_response = await self._generate_final_response(state) + + return self.response_converter.convert_final_response(final_response) + def _build_graph(self) -> CompiledGraph: # Define a new graph. workflow = StateGraph(SupervisorState) @@ -242,7 +260,7 @@ def _build_graph(self) -> CompiledGraph: # Define the nodes of the graph. workflow.add_node(PLANNER, self._plan) workflow.add_node(ROUTER, self._route) - workflow.add_node(FINALIZER, self._generate_final_response) + workflow.add_node(FINALIZER, self._get_converted_final_response) # Set the entrypoint: ENTRY --> (planner | router | finalizer) workflow.add_conditional_edges( diff --git a/src/agents/supervisor/prompts.py b/src/agents/supervisor/prompts.py index a8afc668..2d101616 100644 --- a/src/agents/supervisor/prompts.py +++ b/src/agents/supervisor/prompts.py @@ -60,6 +60,7 @@ - Include ALL the provided code blocks (YAML, JavaScript, JSON, etc.) in the final response. - Remove any information regarding the agents and your decision-making process from your final response. - Do not add any more headers or sub-headers to the final response. +- If there is any YAML config , put the config in or block based on whether it is for new deployment or updating existing deployment. """ FINALIZER_PROMPT_FOLLOW_UP = """ diff --git a/tests/unit/agents/common/test_response_converter.py b/tests/unit/agents/common/test_response_converter.py new file mode 100644 index 00000000..d837151b --- /dev/null +++ b/tests/unit/agents/common/test_response_converter.py @@ -0,0 +1,480 @@ +import pytest +from langchain_core.messages import AIMessage + +from agents.common.constants import FINALIZER, MESSAGES, NEW_YAML, UPDATE_YAML +from agents.common.response_converter import ResponseConverter + + +@pytest.fixture +def response_converter(): + return ResponseConverter() + + +@pytest.mark.parametrize( + "description,yaml_content,expected_results", + [ + ( + "Response with new and update yaml configs", + """ +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: default +spec: + replicas: 3 +``` + + +some text +some text + + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: default +spec: + replicas: 5 +``` +""", + ( + [ + "```yaml\napiVersion: apps/v1\nkind: Deployment\nmetadata:\n name: nginx-deployment\n namespace: default\nspec:\n replicas: 3\n```" + ], + [ + "```yaml\napiVersion: apps/v1\nkind: Deployment\nmetadata:\n name: nginx-deployment\n namespace: default\nspec:\n replicas: 5\n```" + ], + ), + ), + ("Response with no yaml configs", "No YAML content", ([], [])), + ( + "Response with invalid yaml configs", + "\ninvalid yaml\n", + (["invalid yaml"], []), + ), + ( + "Response with two yaml", + """To create an Nginx deployment with 3 replicas, you can use the following YAML configuration: + + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: nginx-oom +spec: + replicas: 3 + selector: + matchLabels: + app: nginx + template: + metadata: + labels: + app: nginx + spec: + containers: + - name: nginx + image: nginx:latest + ports: + - containerPort: 80 +``` + + +To apply this configuration, save it to a file (e.g., `nginx-deployment.yaml`) and run the following command: + +```bash +kubectl apply -f nginx-deployment.yaml +``` + +If you need to update the configuration of the existing Nginx deployment, for example, to change the image version or add environment variables, you can use the following YAML: + + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: nginx-oom +spec: + replicas: 3 + selector: + matchLabels: + app: nginx + template: + metadata: + labels: + app: nginx + spec: + containers: + - name: nginx + image: nginx:1.21.0 # Updated image version + ports: + - containerPort: 80 + env: + - name: NGINX_ENV + value: "production" # New environment variable +``` + + +To apply this update, save the YAML to a file (e.g., `nginx-deployment-update.yaml`) and run: + +```bash +kubectl apply -f nginx-deployment-update.yaml +``` + +This will update the existing deployment with the new configuration.""", + ( + [ + """```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: nginx-oom +spec: + replicas: 3 + selector: + matchLabels: + app: nginx + template: + metadata: + labels: + app: nginx + spec: + containers: + - name: nginx + image: nginx:latest + ports: + - containerPort: 80 +```""" + ], + [ + """```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: nginx-oom +spec: + replicas: 3 + selector: + matchLabels: + app: nginx + template: + metadata: + labels: + app: nginx + spec: + containers: + - name: nginx + image: nginx:1.21.0 # Updated image version + ports: + - containerPort: 80 + env: + - name: NGINX_ENV + value: "production" # New environment variable +```""" + ], + ), + ), + ], +) +def test_extract_yaml(response_converter, description, yaml_content, expected_results): + new_yaml, update_yaml = response_converter._extract_yaml(yaml_content) + assert (new_yaml, update_yaml) == expected_results + + +@pytest.mark.parametrize( + "description, yaml_content,expected_result", + [ + ( + "yaml with yaml marker", + """```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: test + namespace: default +```""", + { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": "test", "namespace": "default"}, + }, + ), + ( + "yaml without yaml marker", + """apiVersion: apps/v1 +kind: Deployment +metadata: + name: test + namespace: default""", + { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": "test", "namespace": "default"}, + }, + ), + ("invalid yaml", "invalid: :\nyaml: content:", None), + ("No yaml", "", None), + ], +) +def test_parse_yamls(response_converter, description, yaml_content, expected_result): + assert response_converter._parse_yamls(yaml_content) == expected_result + + +@pytest.mark.parametrize( + "description, yaml_config,link_type,expected_link", + [ + ( + "Generate link for new deployment", + { + "metadata": {"namespace": "test-ns", "name": "test-deploy"}, + "kind": "Deployment", + }, + NEW_YAML, + "/namespaces/test-ns/Deployment", + ), + ( + "Generate link for updating old deployment", + { + "metadata": {"namespace": "test-ns", "name": "test-deploy"}, + "kind": "Deployment", + }, + UPDATE_YAML, + "/namespaces/test-ns/Deployment/test-deploy", + ), + ( + "Test no link generation", + {"metadata": {"namespace": "test-ns"}, "kind": "Deployment"}, + NEW_YAML, + None, + ), + ], +) +def test_generate_resource_link( + response_converter, description, yaml_config, link_type, expected_link +): + assert ( + response_converter._generate_resource_link(yaml_config, link_type) + == expected_link + ) + + +@pytest.mark.parametrize( + "description, yaml_content,resource_link,link_type,expected_contents", + [ + ( + "Test html generation", + "kind: Deployment", + "/test/link", + NEW_YAML, + [ + 'class="yaml-block', + 'class="yaml"', + 'class="link"', + "/test/link", + "kind: Deployment", + ], + ) + ], +) +def test_create_html_nested_yaml( + response_converter, + description, + yaml_content, + resource_link, + link_type, + expected_contents, +): + html = response_converter._create_html_nested_yaml( + yaml_content, resource_link, link_type + ) + for content in expected_contents: + assert content in html + + +@pytest.mark.parametrize( + "finalizer_response,replacement_html_list,yaml_type,expected", + [ + ( + """Some text before + + metadata: + name: test + + Some text after""", + ["
HTML Replace 1
"], + NEW_YAML, + """Some text before +
HTML Replace 1
+ Some text after""", + ), + ( + """First block + + block1 + + Second block + + block2 +""", + ["
Replace 1
", "
Replace 2
"], + UPDATE_YAML, + """First block +
Replace 1
+ Second block +
Replace 2
""", + ), + ( + """No YAML blocks here""", + ["
HTML
"], + NEW_YAML, + """No YAML blocks here""", + ), + ( + """block""", + [], + NEW_YAML, + """block""", + ), + ], + ids=[ + "single_replacement", + "multiple_replacements", + "no_blocks", + "empty_replacement_list", + ], +) +def test_replace_yaml_with_html( + response_converter, finalizer_response, replacement_html_list, yaml_type, expected +): + result = response_converter._replace_yaml_with_html( + finalizer_response, replacement_html_list.copy(), yaml_type + ) + # Normalize whitespace for comparison + assert " ".join(result.split()) == " ".join(expected.split()) + + +@pytest.mark.parametrize( + "yaml_list,yaml_type,expected", + [ + ( + [ + """ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: nginx-oom""" + ], + NEW_YAML, + [ + f""" +
+apiVersion: apps/v1 +kind: Deployment +metadata: + name: nginx-deployment + namespace: nginx-oom +
+ +
+ [Apply](/namespaces/nginx-oom/Deployment) +
+ + """ + ], + ), + ( + [ + """invalid: :""", + """ +apiVersion: apps/v1 +kind: Service +metadata: + name: test-svc + namespace: test-ns""", + ], + UPDATE_YAML, + [ + """invalid: :""", + f""" +
+apiVersion: apps/v1 +kind: Service +metadata: + name: test-svc + namespace: test-ns +
+ +
+ [Apply](/namespaces/test-ns/Service/test-svc) +
+ + """, + ], + ), + ([], NEW_YAML, []), + ], + ids=["single_valid_yaml", "mixed_valid_invalid", "empty_list"], +) +def test_create_replacement_list(response_converter, yaml_list, yaml_type, expected): + result = response_converter._create_replacement_list(yaml_list, yaml_type) + + # Compare lengths + assert len(result) == len(expected) + + # Compare each element after normalizing whitespace + for res, exp in zip(result, expected, strict=False): + assert " ".join(res.split()) == " ".join(exp.split()) + + +@pytest.mark.parametrize( + "state_content,expected_content", + [ + ( + """ +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: test + namespace: default +``` +""", + """ +
+ ```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: test + namespace: default +``` +
+
+ [Apply](/namespaces/default/Deployment) +
+ + """, + ), + ("No YAML content", "No YAML content"), + ("", ""), + ], + ids=["single_valid_yaml", "No YAML in Content", "Empty Response"], +) +def test_convert_final_response(response_converter, state_content, expected_content): + state = {"messages": [AIMessage(content=state_content, name=FINALIZER)]} + result = response_converter.convert_final_response(state) + assert " ".join(result[MESSAGES][0].content.split()) == " ".join( + expected_content.split() + ) diff --git a/tests/unit/agents/test_graph.py b/tests/unit/agents/test_graph.py index 807403c5..47aa7800 100644 --- a/tests/unit/agents/test_graph.py +++ b/tests/unit/agents/test_graph.py @@ -383,7 +383,8 @@ def test_agent_initialization(self, companion_graph, mock_models, mock_memory): # Verify SupervisorAgent was constructed with correct arguments mock_supervisor_cls.assert_called_once_with( - mock_models, members=["KymaAgent", "KubernetesAgent", "Common"] + mock_models, + members=["KymaAgent", "KubernetesAgent", "Common"], ) mock_build_graph.assert_called_once()