Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: added response conversion node #318

Merged
merged 11 commits into from
Jan 24, 2025
6 changes: 6 additions & 0 deletions src/agents/common/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,12 @@

KYMA_AGENT = "KymaAgent"

RESPONSE_CONVERTER = "ResponseConverter"

NEW_YAML = "New"

UPDATE_YAML = "Update"

SUCCESS_CODE = 200

ERROR_RATE_LIMIT_CODE = 429
272 changes: 272 additions & 0 deletions src/agents/common/response_converter.py
Original file line number Diff line number Diff line change
@@ -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:
tanweersalah marked this conversation as resolved.
Show resolved Hide resolved
"""
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"<YAML-NEW>\n(.*?)\n</YAML-NEW>"
self.update_yaml_pattern = r"<YAML-UPDATE>\n(.*?)\n</YAML-UPDATE>"

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
Comment on lines +65 to +95
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's use TypedDict or Pydantic to parse:

class YAMLMetadata(TypedDict):
    namespace: str
    name: str

class YAMLConfig(TypedDict):
    metadata: YAMLMetadata
    kind: str

def _parse_yamls(self, yaml_config: str) -> Optional[YAMLConfig]:
    """
    Parse YAML string into Python object with error handling.
    Returns None if parsing fails.

    Args:
        yaml_config: YAML configuration string

    Returns:
        Parsed YAML object or None if parsing fails
    
    Example:
        >>> result = converter._parse_yamls("```yaml\\nkind: Deployment\\n```")
        >>> if result:
        >>>     print(f"Parsed YAML: {result}")
        >>> else:
        >>>     print("Failed to parse YAML")
    """
    try:
        if yaml_config[:7] == "```yaml":
            parsed_yaml = yaml.safe_load(yaml_config[8:-4])
        else:
            parsed_yaml = yaml.safe_load(yaml_config)
        
        return parsed_yaml

    except Exception as e:
        logger.error(
            f"Error while parsing the yaml : {yaml_config} , Exception : {e}"
        )
        return None

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As I see if a yaml is generated but doesn't fit the schema, it returns None.
Have you observed the cases where partial yaml resources are generated? We can discuss this.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As agreed will skip this


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"""
<div class="yaml-block>
<div class="yaml">
{yaml_config}
</div>
<div class="link" link-type="{link_type}">
[Apply]({resource_link})
</div>
</div>
"""

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,
}
22 changes: 20 additions & 2 deletions src/agents/supervisor/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -235,14 +244,23 @@ 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)

# 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(
Expand Down
1 change: 1 addition & 0 deletions src/agents/supervisor/prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 <YAML-NEW> </YAML-NEW> or <YAML-UPDATE> </YAML-UPDATE> block based on whether it is for new deployment or updating existing deployment.
"""

FINALIZER_PROMPT_FOLLOW_UP = """
Expand Down
Loading
Loading