Demo Walk the Policy Auditor Agent and System Prompt

# structure_auditor.py
import asyncio
from pathlib import Path
from mcp_client import run_agent

TAGGING_RULES_PATH = Path(__file__).parents[1] / "config" / "tagging-rules.yaml"

PROMPT = """
You are auditing Kubernetes resource labels for FinOps governance through MCP tools.

Use the available Kubernetes MCP tools to:
1. List all namespaces.
2. For each application namespace (skip only namespaces listed in excluded_namespaces;
   also skip kube-* namespaces), fetch deployments, configmaps, and persistent volume claims.
3. Report each resource's name, namespace, kind, and labels exactly as returned by the tools.
4. Evaluate each resource against the tagging rules provided separately.
5. Flag resources that violate required tags or label mappings.
6. Summarize the full label audit in plain English, grouped by namespace. Cover every
   audited namespace and resource — do not truncate mid-report.

Evaluation rules (follow strictly; do not invent policy):
- A required tag is present if ANY key listed under that tag in label_mappings appears
  on the resource (for example, `app` satisfies `application`).
- Do not call a mapped key non-standard, legacy, or invalid when it is in label_mappings.
- Use only required_tags, label_mappings, cost_centers, environments, tiers,
  compliance_levels, resource_types, and excluded_namespaces from the rules.
- Do not invent exclusions, bonus labels, recommended tags, or extra commentary
  beyond what the rules define.

Use only data returned by the tools. Do not invent labels or resources.
"""

def load_tagging_rules() -> str:
    if not TAGGING_RULES_PATH.exists():
        raise FileNotFoundError(f"Tagging rules not found: {TAGGING_RULES_PATH}")
    return TAGGING_RULES_PATH.read_text()

async def main() -> None:
    # Call the shared run_agent function and print the result.
    print(
        await run_agent(
            PROMPT,
            tagging_rules=load_tagging_rules(),
            max_tokens=4096,
        )
    )

if __name__ == "__main__":
    asyncio.run(main())

# mcp_client.py
import os
from typing import Tuple
from langchain.chat_models import ChatOpenAI
from langchain.schema import SystemMessage, HumanMessage
# The following conversion and agent creation helpers are assumed to be implemented
# elsewhere in the codebase:
# - convert_mcp_to_langchain_tools
async def get_mcp_tools():
    # convert_mcp_to_langchain_tools should return (tools, cleanup_coroutine)
    return await convert_mcp_to_langchain_tools(MCP_SERVERS)

def build_llm(max_tokens: int | None = None) -> ChatOpenAI:
    return ChatOpenAI(
        model_name=os.environ["OPENAI_MODEL_ID"],
        openai_api_base=os.environ.get("OPENAI_BASE_URL"),
        openai_api_key=os.environ.get("OPENAI_API_KEY"),
        temperature=float(os.getenv("OPENAI_TEMPERATURE", "0.3")),
        max_tokens=max_tokens or int(os.getenv("OPENAI_MAX_TOKENS", "512")),
    )

async def run_agent(
    prompt: str,
    max_tokens: int | None = None,
    tagging_rules: str | None = None,
) -> str:
    tools, cleanup = await get_mcp_tools()
    try:
        agent = create_agent(build_llm(max_tokens=max_tokens), tools)
        messages = []
        if tagging_rules:
            # Strict system-level guidance that the model must follow.
            messages.append(
                SystemMessage(
                    content=(
                        "Apply the tagging rules below when evaluating cluster resources. "
                        "Treat label_mappings as the only accepted keys for each required tag. "
                        "Do not invent extra policy, exclusions, or judgments beyond this YAML.\n\n"
                        f"Tagging rules (YAML):\n{tagging_rules}"
                    )
                )
            )
        messages.append(HumanMessage(content=prompt))
        result = await agent.arun({"messages": messages})
        # The agent may return a string or a structured message list; handle both.
        if isinstance(result, dict) and "messages" in result:
            return result["messages"][-1].content
        return result
    finally:
        await cleanup()