Friday, September 11, 2026

BUILDING AN LLM-BASED AGENTIC AI FOR KUBERNETES DEPLOYMENT AUTOMATION

 


INTRODUCTION AND VISION

Greetings, Siemens colleague! Today we embark on an extraordinary journey into the realm of intelligent automation for cloud-native infrastructure. We will explore how to build a sophisticated LLM-based Agentic AI system that transforms the way we interact with Kubernetes, Docker, and Go applications. This is not merely a code generator; it is an intelligent assistant that understands context, applies best practices, validates every decision, and produces production-ready artifacts.

The system we are about to design represents a paradigm shift in infrastructure management. Imagine describing your deployment needs in natural language and receiving not just YAML files, but a complete, validated, production-ready deployment with security configurations, monitoring integration, ingress controllers, custom resource definitions, operators, comprehensive documentation, and build scripts. This Agentic AI will handle the entire lifecycle: creating new deployments, evolving existing ones, migrating from basic deployment descriptors to sophisticated Helm Charts, and ensuring every generated artifact meets the highest standards of correctness and reliability.

The foundation of this system rests on several pillars: a multi-agent architecture where specialized agents collaborate under the guidance of a coordinator, advanced knowledge retrieval using RAG and GraphRAG to maintain up-to-date expertise, a comprehensive toolset for code generation and validation, and support for diverse hardware platforms including Intel, AMD ROCm, Apple MPS, and Nvidia CUDA.

ARCHITECTURAL FOUNDATIONS

THE MULTI-AGENT PARADIGM

The heart of our Agentic AI is a carefully orchestrated multi-agent system. This architecture provides clear separation of concerns, allowing each agent to develop deep expertise in its domain while the coordinator ensures harmonious collaboration. The multi-agent approach offers several critical advantages over monolithic systems: it enables parallel processing of complex tasks, provides modularity for easier maintenance and extension, allows specialized optimization for each domain, and creates natural boundaries for testing and validation.

The architecture follows a hub-and-spoke model where the Coordinator Agent sits at the center, managing communication and workflow between specialized agents. Each specialized agent possesses deep knowledge in its domain, maintained through continuous learning from documentation, code repositories, and validated examples. The agents communicate through a structured message protocol, passing context, requirements, and artifacts between them.

THE COORDINATOR AGENT

The Coordinator Agent serves as the intelligent orchestrator of the entire system. When a user makes a request, the Coordinator performs several critical functions. First, it analyzes the natural language input to extract intent, requirements, and constraints. It identifies which specialized agents need to be involved and in what sequence. It maintains the overall context of the conversation, remembering previous decisions and artifacts. It manages the workflow, ensuring dependencies are respected and that agents receive the information they need. Finally, it synthesizes the outputs from multiple agents into a coherent solution.

The Coordinator uses a sophisticated planning mechanism. When it receives a request like "Create a Go microservice with PostgreSQL, expose it via Ingress with TLS, and add Prometheus monitoring," it constructs an execution plan. This plan might look like the following sequence: First, the Go Application Agent designs the service structure. Then, the Docker Agent creates optimized container images. Next, the Kubernetes Core Agent generates base deployment and service manifests. The Database Agent configures PostgreSQL with persistent storage. The Ingress Agent sets up routing and TLS certificates. The Monitoring Agent integrates Prometheus scraping. Finally, the Validation Agent verifies all configurations before presenting the complete solution to the user.

SPECIALIZED EXPERT AGENTS

Each specialized agent focuses on a specific domain within the Kubernetes ecosystem. Let us explore the key agents and their responsibilities.

The Docker Agent possesses comprehensive knowledge of containerization best practices. It generates Dockerfiles optimized for different languages and frameworks, with particular expertise in Go applications. It understands multi-stage builds to minimize image size, security scanning and vulnerability management, layer caching strategies for faster builds, and platform-specific optimizations for different CPU architectures. When creating a Dockerfile for a Go application, it ensures proper dependency management, minimal base images like Alpine or distroless, appropriate user permissions avoiding root execution, and health check configurations.

The Kubernetes Core Agent handles the fundamental Kubernetes resources. It generates Deployments with proper replica management, resource requests and limits, health checks and readiness probes, and update strategies. It creates Services for internal and external communication, ConfigMaps and Secrets for configuration management, PersistentVolumeClaims for stateful applications, and ServiceAccounts with appropriate RBAC permissions. This agent ensures all manifests follow Kubernetes best practices and include proper labels, annotations, and selectors.

The Ingress Agent specializes in external access configuration. It generates Ingress resources for HTTP and HTTPS routing, configures TLS certificates using cert-manager or manual secrets, sets up path-based and host-based routing rules, and integrates with various Ingress controllers like Nginx, Traefik, or Istio. It understands rate limiting, authentication, and other advanced ingress features.

The CRD and Operator Agent handles Kubernetes extensibility. It generates Custom Resource Definitions with proper validation schemas, creates Operators using the Operator SDK or Kubebuilder, implements reconciliation loops in Go, and handles complex state management. This agent understands the intricacies of extending Kubernetes and ensures custom resources integrate seamlessly with the existing ecosystem.

The Monitoring Agent integrates observability solutions. It configures Prometheus ServiceMonitors and PodMonitors, sets up Grafana dashboards, integrates distributed tracing with Jaeger or Zipkin, and configures log aggregation with the ELK stack or Loki. It ensures comprehensive visibility into application and infrastructure health.

The Security Agent focuses on hardening deployments. It configures Pod Security Policies or Pod Security Standards, sets up Network Policies for traffic control, manages secrets encryption and rotation, implements security contexts with appropriate capabilities, and integrates vulnerability scanning into the deployment pipeline.

The Helm Agent specializes in package management. It converts raw Kubernetes manifests into Helm Charts with proper templating, creates values files for different environments, implements chart dependencies and hooks, and follows Helm best practices for versioning and distribution.

The Validation Agent serves as the quality gatekeeper. It performs static analysis of all generated YAML and code, validates Kubernetes manifests against API schemas, checks for security vulnerabilities and misconfigurations, verifies resource dependencies and references, and ensures compliance with organizational policies. This agent prevents the system from producing incorrect or unsafe configurations.

KNOWLEDGE MANAGEMENT THROUGH RAG AND GRAPHRAG

The intelligence of our Agentic AI depends critically on its knowledge base. Unlike static systems, our AI continuously learns from authoritative sources, maintaining current and accurate information about Kubernetes, Docker, Go, and related technologies.

RETRIEVAL AUGMENTED GENERATION

RAG forms the foundation of the system's knowledge retrieval mechanism. The process begins with document ingestion, where the system automatically downloads and processes documentation from official sources like the Kubernetes documentation repository, Docker documentation, Go language specifications, Helm documentation, and operator framework guides. These documents undergo sophisticated processing.

The ingestion pipeline chunks documents into semantically meaningful segments, typically paragraphs or logical sections. Each chunk is then converted into high-dimensional vector embeddings using specialized models. These embeddings capture the semantic meaning of the text, allowing the system to find relevant information based on conceptual similarity rather than just keyword matching. The embeddings are stored in a vector database like Chroma, Pinecone, or Weaviate, enabling efficient similarity search.

When an agent needs information to complete a task, it formulates a query based on the current context. For example, if the Docker Agent needs to create a multi-stage Dockerfile for a Go application, it might query for "Go multi-stage Docker build best practices" or "optimizing Go Docker images." The RAG system converts this query into an embedding and searches the vector database for the most similar document chunks. These retrieved chunks provide the agent with current, authoritative information to inform its decisions.

The RAG system includes a freshness mechanism. It periodically checks for updates to source documentation and re-ingests changed content. This ensures agents always work with the latest best practices and API specifications. When Kubernetes releases a new version with API changes, the RAG system automatically incorporates this information, preventing the generation of deprecated configurations.

GRAPHRAG FOR RELATIONSHIP MODELING

While traditional RAG excels at retrieving relevant text chunks, GraphRAG adds a crucial dimension: understanding relationships between entities. In the Kubernetes ecosystem, relationships are fundamental. A Deployment references a ConfigMap, which might be used by multiple Services, which are exposed through an Ingress, monitored by a ServiceMonitor, and protected by a NetworkPolicy. Understanding these relationships is essential for generating coherent, complete deployments.

GraphRAG builds a knowledge graph where nodes represent entities like Kubernetes resources, Docker images, Go packages, configuration parameters, and best practices. Edges represent relationships such as "depends on," "configures," "monitors," "secures," and "exposes." This graph is constructed through several mechanisms.

First, the system extracts entities and relationships from documentation using natural language processing. When processing a document about Deployments, it identifies that Deployments can reference ConfigMaps and Secrets, that they create ReplicaSets, and that they are exposed by Services. Second, the system learns from validated examples in code repositories. By analyzing real-world Kubernetes manifests, it discovers common patterns and relationships. Third, the system updates the graph based on successful deployments, reinforcing patterns that work well together.

When an agent needs to make a decision, it queries both the vector database for textual information and the knowledge graph for relationship context. For instance, if the Coordinator determines that a deployment needs monitoring, it queries the graph to understand that a Deployment should have a Service with specific labels, and that Service should be targeted by a ServiceMonitor with matching selectors. This relationship awareness prevents common configuration errors and ensures all components integrate correctly.

The knowledge graph also enables powerful reasoning capabilities. If a user asks to add authentication to an existing deployment, the system can traverse the graph to understand what components need modification: the Ingress needs authentication annotations, the application might need an authentication library, environment variables might need to be added to the Deployment, and a Secret might need to be created for credentials. This holistic understanding distinguishes our Agentic AI from simple code generators.

TOOL ECOSYSTEM FOR COMPREHENSIVE AUTOMATION

The Agentic AI requires a rich set of tools to interact with the external world, validate its outputs, and manage complex workflows. These tools extend the capabilities of the LLM-based agents, allowing them to perform concrete actions rather than just generating text.

CODE REPOSITORY BROWSING AND ANALYSIS

The system includes sophisticated tools for interacting with code repositories. The Repository Browser tool can clone Git repositories, navigate directory structures, read file contents, and analyze code patterns. This capability is essential for several use cases.

When a user asks to evolve an existing deployment, the system first browses the current repository to understand the existing structure. It identifies which Kubernetes manifests are present, what Docker images are being used, what configuration files exist, and how components are organized. This analysis informs the evolution strategy, ensuring changes integrate smoothly with existing code.

The Repository Analyzer tool performs deeper inspection. It can parse Kubernetes YAML to extract resource definitions and relationships, analyze Dockerfiles to understand build processes and dependencies, examine Go code to identify service endpoints and configuration requirements, and detect patterns that suggest best practices or anti-patterns. This analysis feeds into the knowledge graph, allowing the system to learn from real-world examples.

WEB SEARCH FOR CURRENT INFORMATION

While the RAG system maintains a comprehensive knowledge base, some information requires real-time retrieval. The Web Search tool allows agents to find current information about newly released features, community best practices, security vulnerabilities, and compatibility matrices. This tool is particularly valuable for the Validation Agent, which can search for known issues with specific Kubernetes versions or Docker base images before approving a configuration.

FILE SYSTEM OPERATIONS

The File System tool enables the Agentic AI to create the complete directory structure for a project. When generating a new deployment, the system creates organized directories for Kubernetes manifests, Docker configurations, Go source code, Helm charts, documentation, and build scripts. It writes all generated files to appropriate locations, ensuring a clean, professional project structure.

The tool follows conventions for each type of project. For a Go application with Kubernetes deployment, it might create a structure like this:

project-root/
  cmd/
    service-name/
      main.go
  internal/
    handlers/
    models/
    config/
  deployments/
    kubernetes/
      base/
        deployment.yaml
        service.yaml
        configmap.yaml
      overlays/
        dev/
        staging/
        production/
    docker/
      Dockerfile
      .dockerignore
  charts/
    service-name/
      Chart.yaml
      values.yaml
      templates/
  scripts/
    build.sh
    deploy.sh
    validate.sh
  docs/
    README.md
    ARCHITECTURE.md
  Makefile
  go.mod
  go.sum

This structure provides clear organization, separates concerns, and follows community conventions.

VALIDATION TOOLS

The Validation toolset is perhaps the most critical component ensuring the system never produces incorrect configurations. These tools include several specialized validators.

The Kubernetes Manifest Validator uses tools like kubeval or kubectl dry-run to verify that all YAML manifests are syntactically correct and semantically valid against the target Kubernetes API version. It checks that all required fields are present, that field values are of the correct type, that resource references are valid, and that selectors match labels correctly.

The Docker Validator analyzes Dockerfiles for security issues, inefficient layer construction, missing health checks, and improper use of base images. It can run tools like hadolint to enforce Dockerfile best practices.

The Go Code Validator runs static analysis tools like go vet, golint, and staticcheck on generated Go code. It ensures the code compiles, follows Go conventions, handles errors appropriately, and avoids common pitfalls.

The Security Validator scans for security issues using tools like trivy for container vulnerabilities, kube-score for Kubernetes best practices, and custom rules for organizational security policies. It ensures no secrets are hardcoded, that containers run as non-root users, that security contexts are properly configured, and that network policies restrict traffic appropriately.

The Integration Validator performs end-to-end testing by deploying configurations to a test Kubernetes cluster, verifying that all resources are created successfully, checking that services are accessible, running health checks, and validating monitoring integration. This final validation step ensures the entire system works together correctly before presenting the solution to the user.

BUILD AND DEPLOYMENT AUTOMATION

The system generates comprehensive build and deployment automation. The Make Script Generator creates Makefiles with targets for building Docker images, running tests, deploying to different environments, and cleaning up resources. The CI/CD Integration tool generates pipeline configurations for systems like Jenkins, GitLab CI, GitHub Actions, or Azure DevOps, automating the entire deployment workflow.

LLM INFRASTRUCTURE AND HARDWARE SUPPORT

The Agentic AI must support diverse deployment scenarios, from local development on personal workstations to large-scale production deployments on cloud infrastructure. This requires flexible LLM infrastructure supporting various hardware platforms.

LOCAL LLM DEPLOYMENT

For local development and testing, the system supports running LLMs directly on developer workstations. This provides fast iteration, data privacy, and offline capability. The system integrates with frameworks like llama.cpp, Ollama, and LocalAI that enable efficient local inference.

The hardware abstraction layer detects the available compute resources and selects appropriate optimizations. On Apple Silicon Macs, it uses Metal Performance Shaders for GPU acceleration. On systems with Nvidia GPUs, it leverages CUDA for maximum performance. For AMD GPUs, it uses ROCm to access GPU compute capabilities. On Intel systems without discrete GPUs, it optimizes for CPU inference using AVX2 and AVX-512 instructions.

The model selection mechanism chooses appropriate model sizes based on available resources. On a laptop with 16GB of RAM, it might use a 7B parameter model. On a workstation with 64GB and a powerful GPU, it can use larger 13B or 33B models for improved reasoning. The system includes a model registry mapping tasks to optimal models, balancing capability with resource constraints.

REMOTE LLM SERVICES

For production deployments, the system supports remote LLM services offering greater scale and capability. It integrates with OpenAI's API for GPT-4 and GPT-3.5, Anthropic's Claude for complex reasoning tasks, Google's PaLM for specialized applications, and Azure OpenAI for enterprise deployments. The abstraction layer provides a unified interface, allowing agents to work with any backend without code changes.

The remote service integration includes sophisticated error handling and fallback mechanisms. If a primary service is unavailable, the system automatically fails over to alternative providers. It implements rate limiting and request queuing to respect API quotas. It caches responses for identical queries to reduce costs and latency.

HYBRID DEPLOYMENT ARCHITECTURE

The most sophisticated deployments use a hybrid approach, combining local and remote LLMs strategically. Simple tasks like parsing user input or formatting output run on local models for low latency and cost. Complex reasoning tasks like planning deployment strategies or analyzing security implications use powerful remote models. The routing logic considers task complexity, latency requirements, cost constraints, and data sensitivity when deciding where to execute each inference.

The system includes a performance monitoring component that tracks inference latency, accuracy, and cost for each model and task type. This data feeds into a continuous optimization process that refines the routing decisions over time, learning which models perform best for which tasks.

GPU ACCELERATION DETAILS

Supporting diverse GPU architectures requires careful abstraction and optimization. The system includes platform-specific code paths for each GPU vendor.

For Nvidia CUDA, the system uses PyTorch or TensorFlow with CUDA support, cuBLAS for optimized matrix operations, and TensorRT for inference optimization. It detects the GPU compute capability and selects appropriate kernel implementations. On newer Ampere or Ada Lovelace GPUs, it leverages tensor cores for mixed-precision inference, achieving significant speedups.

For AMD ROCm, the system uses ROCm-enabled PyTorch, MIOpen for GPU-accelerated operations, and platform-specific optimizations for RDNA or CDNA architectures. While ROCm support is less mature than CUDA, the system includes workarounds for known issues and fallback paths when specific operations are not well-supported.

For Apple Metal, the system uses Core ML for model inference, Metal Performance Shaders for GPU compute, and the Accelerate framework for CPU optimizations. On Apple Silicon, the unified memory architecture allows efficient data sharing between CPU and GPU, which the system exploits for reduced memory overhead.

For Intel GPUs, the system uses oneAPI and SYCL for GPU acceleration, oneDNN for deep learning primitives, and OpenVINO for optimized inference. Intel's discrete GPUs like Arc are increasingly capable for AI workloads, and the system takes advantage of their XMX engines for matrix operations.

The abstraction layer presents a unified interface to the agents, hiding the complexity of platform-specific optimizations. An agent simply requests inference on a particular model, and the infrastructure handles device selection, memory management, and execution.

WORKFLOW AND USER INTERACTION

Understanding how the system operates end-to-end illuminates the power of the multi-agent architecture. Let us walk through a comprehensive example.

INITIAL USER REQUEST

A user interacts with the Coordinator Agent through a natural language interface. The user might say: "I need to deploy a Go-based REST API that connects to PostgreSQL. The API should be accessible externally with HTTPS, monitored with Prometheus, and deployed using Helm. I want separate configurations for development and production environments."

The Coordinator Agent parses this request, identifying several key requirements: a Go application requiring development, a PostgreSQL database requiring configuration, Docker containerization for the Go application, Kubernetes deployment manifests, an Ingress with TLS for external access, Prometheus monitoring integration, Helm chart creation for package management, and environment-specific configurations.

PLANNING AND TASK DECOMPOSITION

The Coordinator constructs a detailed execution plan. It determines the sequence of operations, identifies dependencies between tasks, and assigns tasks to specialized agents. The plan might look like this:

First, the Go Application Agent will design the REST API structure, including endpoint definitions, database models, configuration management, and error handling. Second, the Database Agent will create PostgreSQL configuration, including deployment manifests, persistent volume claims, connection secrets, and initialization scripts. Third, the Docker Agent will generate an optimized Dockerfile with multi-stage builds, minimal base images, and proper security configurations. Fourth, the Kubernetes Core Agent will create base deployment and service manifests with health checks, resource limits, and proper labels. Fifth, the Ingress Agent will configure external access with TLS certificate management, path routing, and rate limiting. Sixth, the Monitoring Agent will set up Prometheus integration with ServiceMonitor configuration, custom metrics endpoints, and Grafana dashboard templates. Seventh, the Helm Agent will convert all manifests to a Helm chart with templating for environment-specific values, chart dependencies, and deployment hooks. Finally, the Validation Agent will verify all configurations through static analysis, security scanning, and integration testing.

The Coordinator also identifies that some tasks can run in parallel. For example, the Go application development and PostgreSQL configuration can proceed simultaneously, as they have minimal dependencies initially.

AGENT EXECUTION AND COLLABORATION

Each agent executes its assigned tasks, leveraging the RAG system for knowledge and tools for concrete actions.

The Go Application Agent begins by querying the RAG system for best practices in Go REST API development. It retrieves information about popular frameworks like Gin or Echo, database connection patterns using sqlx or GORM, configuration management with Viper, and structured logging with zap or logrus. Using this knowledge, it generates a complete Go application structure.

The generated main.go file includes proper initialization, graceful shutdown handling, and configuration loading. The handlers package contains endpoint implementations with proper error handling and input validation. The models package defines database schemas using struct tags for ORM mapping. The config package handles environment-specific configuration with sensible defaults.

The Database Agent queries the knowledge graph to understand PostgreSQL deployment patterns in Kubernetes. It discovers that PostgreSQL should use a StatefulSet for stable network identities, requires a PersistentVolumeClaim for data storage, needs a headless Service for pod discovery, and should have initialization ConfigMaps for schema creation. It generates all these resources with proper configurations.

The Docker Agent creates a multi-stage Dockerfile. The first stage uses the official Go image to build the application, copying source code, downloading dependencies, and compiling the binary. The second stage uses a minimal base image like Alpine or distroless, copies only the compiled binary, sets up a non-root user, and configures health check endpoints. The resulting image is small, secure, and efficient.

The Kubernetes Core Agent generates deployment manifests with careful attention to best practices. The Deployment includes replica configuration for high availability, resource requests and limits based on expected load, liveness and readiness probes pointing to health endpoints, environment variables for configuration, and volume mounts for secrets and config maps. The Service exposes the application internally with appropriate selectors.

The Ingress Agent configures external access. It generates an Ingress resource with host-based routing, TLS configuration referencing a cert-manager ClusterIssuer for automatic certificate provisioning, path-based routing to the Service, and annotations for the Nginx Ingress Controller specifying rate limits and CORS policies.

The Monitoring Agent integrates Prometheus monitoring. It adds annotations to the Service for Prometheus scraping, creates a ServiceMonitor custom resource with proper selectors and endpoints, generates a Grafana dashboard JSON with key metrics like request rate, error rate, latency percentiles, and database connection pool status. It also modifies the Go application to expose metrics at the /metrics endpoint using the Prometheus client library.

The Helm Agent takes all the generated Kubernetes manifests and transforms them into a Helm chart. It creates a Chart.yaml with metadata, a values.yaml with configurable parameters like replica count, image repository and tag, resource limits, ingress hostname, and database connection details. It converts the manifests into templates, replacing hardcoded values with template variables. It creates separate values files for development and production environments with appropriate overrides.

VALIDATION AND QUALITY ASSURANCE

Before presenting the solution to the user, the Validation Agent performs comprehensive checks. It runs kubeval against all Kubernetes manifests to ensure API compliance. It uses hadolint on the Dockerfile to check for best practices. It compiles the Go code and runs go vet to catch potential issues. It performs security scanning with trivy on the Docker image to identify vulnerabilities. It checks that all resource references are valid, that selectors match labels, and that dependencies are satisfied.

The Validation Agent also performs integration testing. It deploys the Helm chart to a test Kubernetes cluster, waits for all pods to become ready, runs health checks against the application endpoints, verifies that Prometheus is scraping metrics, and checks that the Ingress is routing traffic correctly. Only after all validations pass does the system consider the solution complete.

SOLUTION PRESENTATION AND DOCUMENTATION

The Coordinator Agent synthesizes all the outputs into a comprehensive solution package. It organizes all files into the proper directory structure. It generates a detailed README.md explaining the project structure, how to build and deploy the application, configuration options, and monitoring setup. It creates a Makefile with targets for common operations like building the Docker image, running tests, deploying to different environments, and cleaning up resources.

The user receives not just code and configuration files, but a complete, production-ready solution with documentation, automation, and confidence that everything has been thoroughly validated.

EVOLUTION AND MIGRATION CAPABILITIES

The Agentic AI excels not only at creating new deployments but also at evolving existing ones and migrating between deployment paradigms.

EVOLVING EXISTING DEPLOYMENTS

When a user requests changes to an existing deployment, the system follows a careful process. First, the Repository Browser tool clones the existing repository and analyzes its structure. The Coordinator identifies what resources currently exist, how they are configured, and what patterns are being used.

If the user asks to add a new feature, like integrating Redis for caching, the system determines what changes are necessary. The Database Agent generates Redis deployment manifests. The Kubernetes Core Agent modifies the application Deployment to add Redis connection configuration. The Go Application Agent updates the application code to use Redis for caching. The Validation Agent ensures the changes integrate correctly with existing components.

The system is careful to preserve existing configurations and patterns. If the current deployment uses specific labels or annotations, the new resources follow the same conventions. If there are custom security policies, the new components comply with them. This consistency is crucial for maintaining a coherent deployment.

MIGRATING TO HELM CHARTS

Many organizations start with simple Kubernetes YAML files and later want to migrate to Helm for better management. The Agentic AI automates this migration.

The Helm Agent analyzes the existing YAML files to identify resources and their relationships. It extracts configurable values, like image tags, replica counts, and resource limits. It creates a Helm chart structure with these values parameterized. It generates appropriate templates with conditional logic for optional features. It creates values files for different environments based on any existing variations.

The migration preserves all existing functionality while adding the benefits of Helm: easier updates, environment-specific configurations, and package management. The Validation Agent ensures the Helm chart produces identical resources to the original YAML files when deployed with default values, guaranteeing a safe migration.

COMPREHENSIVE RUNNING EXAMPLE

To demonstrate the complete capabilities of the Agentic AI system, let us present a full running example. This example will show a production-ready implementation of the core components, not simplified or mocked versions.

EXAMPLE SCENARIO

We will build a complete Agentic AI system that can generate a Go-based microservice with PostgreSQL backend, Docker containerization, Kubernetes deployment with Ingress and monitoring, and Helm chart packaging. The system will include the Coordinator Agent, specialized agents for each domain, RAG integration, validation tools, and support for multiple GPU platforms.

COORDINATOR AGENT IMPLEMENTATION

The Coordinator Agent is implemented in Go, providing a robust foundation for the multi-agent system. Here is the core implementation:

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"
    "sync"
    "time"
)

// Message represents communication between agents
type Message struct {
    From      string                 `json:"from"`
    To        string                 `json:"to"`
    Type      string                 `json:"type"`
    Content   map[string]interface{} `json:"content"`
    Timestamp time.Time              `json:"timestamp"`
}

// Agent interface that all agents must implement
type Agent interface {
    Name() string
    ProcessMessage(ctx context.Context, msg Message) (Message, error)
    Initialize(ctx context.Context, config AgentConfig) error
}

// AgentConfig holds configuration for agents
type AgentConfig struct {
    LLMEndpoint   string
    RAGEndpoint   string
    GraphEndpoint string
    ToolRegistry  *ToolRegistry
    ValidationCfg *ValidationConfig
}

// ValidationConfig holds validation settings
type ValidationConfig struct {
    EnableStaticAnalysis bool
    EnableSecurityScan   bool
    EnableIntegrationTest bool
    TestClusterKubeconfig string
}

// Coordinator manages all agents and orchestrates workflows
type Coordinator struct {
    agents       map[string]Agent
    messageQueue chan Message
    config       AgentConfig
    mu           sync.RWMutex
    context      *ConversationContext
}

// ConversationContext maintains state across interactions
type ConversationContext struct {
    UserID        string
    SessionID     string
    History       []Message
    Artifacts     map[string]interface{}
    mu            sync.RWMutex
}

// NewCoordinator creates a new coordinator instance
func NewCoordinator(config AgentConfig) *Coordinator {
    return &Coordinator{
        agents:       make(map[string]Agent),
        messageQueue: make(chan Message, 100),
        config:       config,
        context: &ConversationContext{
            History:   make([]Message, 0),
            Artifacts: make(map[string]interface{}),
        },
    }
}

// RegisterAgent adds an agent to the coordinator
func (c *Coordinator) RegisterAgent(agent Agent) error {
    c.mu.Lock()
    defer c.mu.Unlock()

    if err := agent.Initialize(context.Background(), c.config); err != nil {
        return fmt.Errorf("failed to initialize agent %s: %w", agent.Name(), err)
    }

    c.agents[agent.Name()] = agent
    log.Printf("Registered agent: %s", agent.Name())
    return nil
}

// ProcessUserRequest handles incoming user requests
func (c *Coordinator) ProcessUserRequest(ctx context.Context, request string) (string, error) {
    log.Printf("Processing user request: %s", request)

    // Parse the user request to understand intent
    intent, err := c.parseIntent(ctx, request)
    if err != nil {
        return "", fmt.Errorf("failed to parse intent: %w", err)
    }
    log.Printf("Parsed intent: %+v", intent)

    // Create execution plan based on intent
    plan, err := c.createExecutionPlan(ctx, intent)
    if err != nil {
        return "", fmt.Errorf("failed to create execution plan: %w", err)
    }
    log.Printf("Created execution plan with %d steps", len(plan.Steps))

    // Execute the plan
    result, err := c.executePlan(ctx, plan)
    if err != nil {
        return "", fmt.Errorf("failed to execute plan: %w", err)
    }
    log.Printf("Plan execution completed successfully")

    // Validate the complete solution
    if err := c.validateSolution(ctx, result); err != nil {
        return "", fmt.Errorf("validation failed: %w", err)
    }
    log.Printf("Solution validation passed")

    // Generate final output with documentation
    output, err := c.generateOutput(ctx, result)
    if err != nil {
        return "", fmt.Errorf("failed to generate output: %w", err)
    }

    return output, nil
}

// Intent represents the parsed user intention
type Intent struct {
    Action       string
    Resources    []string
    Requirements map[string]interface{}
    Constraints  map[string]interface{}
}

// parseIntent uses LLM to understand user request
func (c *Coordinator) parseIntent(ctx context.Context, request string) (*Intent, error) {
    llmClient := NewLLMClient(c.config.LLMEndpoint)

    prompt := fmt.Sprintf(`Analyze the following user request and extract structured intent.
Identify the action (create, update, delete, migrate), resources involved, requirements, and constraints.

User Request: %s

Respond with JSON containing: action, resources, requirements, constraints`, request)

    response, err := llmClient.Generate(ctx, prompt)
    if err != nil {
        return nil, fmt.Errorf("LLM generation failed: %w", err)
    }

    var intent Intent
    if err := json.Unmarshal([]byte(response), &intent); err != nil {
        return nil, fmt.Errorf("failed to parse intent JSON: %w", err)
    }

    return &intent, nil
}

// ExecutionPlan represents the workflow to fulfill user intent
type ExecutionPlan struct {
    Steps        []PlanStep
    Dependencies map[string][]string
}

// PlanStep represents a single step in the execution plan
type PlanStep struct {
    ID          string
    AgentName   string
    Task        string
    Inputs      map[string]interface{}
    DependsOn   []string
}

// createExecutionPlan generates a workflow plan
func (c *Coordinator) createExecutionPlan(ctx context.Context, intent *Intent) (*ExecutionPlan, error) {
    plan := &ExecutionPlan{
        Steps:        make([]PlanStep, 0),
        Dependencies: make(map[string][]string),
    }

    switch intent.Action {
    case "create":
        if contains(intent.Resources, "go-service") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-1",
                AgentName: "go-application-agent",
                Task:      "generate-go-service",
                Inputs:    intent.Requirements,
                DependsOn: []string{},
            })
        }

        if contains(intent.Resources, "postgresql") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-2",
                AgentName: "database-agent",
                Task:      "configure-postgresql",
                Inputs:    intent.Requirements,
                DependsOn: []string{},
            })
        }

        if contains(intent.Resources, "docker") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-3",
                AgentName: "docker-agent",
                Task:      "create-dockerfile",
                Inputs:    intent.Requirements,
                DependsOn: []string{"step-1"},
            })
        }

        if contains(intent.Resources, "kubernetes") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-4",
                AgentName: "kubernetes-core-agent",
                Task:      "generate-k8s-manifests",
                Inputs:    intent.Requirements,
                DependsOn: []string{"step-2", "step-3"},
            })
        }

        if contains(intent.Resources, "ingress") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-5",
                AgentName: "ingress-agent",
                Task:      "configure-ingress",
                Inputs:    intent.Requirements,
                DependsOn: []string{"step-4"},
            })
        }

        if contains(intent.Resources, "monitoring") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-6",
                AgentName: "monitoring-agent",
                Task:      "setup-prometheus",
                Inputs:    intent.Requirements,
                DependsOn: []string{"step-4"},
            })
        }

        if contains(intent.Resources, "helm") {
            plan.Steps = append(plan.Steps, PlanStep{
                ID:        "step-7",
                AgentName: "helm-agent",
                Task:      "create-helm-chart",
                Inputs:    intent.Requirements,
                DependsOn: []string{"step-4", "step-5", "step-6"},
            })
        }

        plan.Steps = append(plan.Steps, PlanStep{
            ID:        "step-8",
            AgentName: "validation-agent",
            Task:      "validate-all",
            Inputs:    intent.Requirements,
            DependsOn: getAllStepIDs(plan.Steps),
        })

    case "migrate":
        plan.Steps = append(plan.Steps, PlanStep{
            ID:        "step-1",
            AgentName: "repository-browser",
            Task:      "analyze-existing",
            Inputs:    intent.Requirements,
            DependsOn: []string{},
        })

        plan.Steps = append(plan.Steps, PlanStep{
            ID:        "step-2",
            AgentName: "helm-agent",
            Task:      "migrate-to-helm",
            Inputs:    intent.Requirements,
            DependsOn: []string{"step-1"},
        })

        plan.Steps = append(plan.Steps, PlanStep{
            ID:        "step-3",
            AgentName: "validation-agent",
            Task:      "validate-migration",
            Inputs:    intent.Requirements,
            DependsOn: []string{"step-2"},
        })

    default:
        return nil, fmt.Errorf("unsupported action: %s", intent.Action)
    }

    return plan, nil
}

// PlanResult holds the results of plan execution
type PlanResult struct {
    Artifacts map[string]interface{}
    Metadata  map[string]interface{}
}

// executePlan runs the execution plan
func (c *Coordinator) executePlan(ctx context.Context, plan *ExecutionPlan) (*PlanResult, error) {
    result := &PlanResult{
        Artifacts: make(map[string]interface{}),
        Metadata:  make(map[string]interface{}),
    }

    completed := make(map[string]bool)
    stepResults := make(map[string]interface{})

    for {
        allDone := true
        progress := false

        for _, step := range plan.Steps {
            if completed[step.ID] {
                continue
            }

            allDone = false

            // Check if dependencies are satisfied
            canExecute := true
            for _, depID := range step.DependsOn {
                if !completed[depID] {
                    canExecute = false
                    break
                }
            }

            if !canExecute {
                continue
            }

            // Execute the step
            log.Printf("Executing step %s: %s on agent %s", step.ID, step.Task, step.AgentName)

            agent, exists := c.agents[step.AgentName]
            if !exists {
                return nil, fmt.Errorf("agent not found: %s", step.AgentName)
            }

            // Prepare message for agent
            msg := Message{
                From:      "coordinator",
                To:        step.AgentName,
                Type:      step.Task,
                Content:   step.Inputs,
                Timestamp: time.Now(),
            }

            // Add results from dependent steps
            for _, depID := range step.DependsOn {
                if depResult, ok := stepResults[depID]; ok {
                    msg.Content["dependency_"+depID] = depResult
                }
            }

            // Process message with agent
            response, err := agent.ProcessMessage(ctx, msg)
            if err != nil {
                return nil, fmt.Errorf("step %s failed: %w", step.ID, err)
            }

            // Store step result
            stepResults[step.ID] = response.Content
            completed[step.ID] = true
            progress = true

            log.Printf("Step %s completed successfully", step.ID)
        }

        if allDone {
            break
        }

        if !progress {
            return nil, fmt.Errorf("execution deadlock: no progress made")
        }
    }

    // Collect all artifacts
    for stepID, stepResult := range stepResults {
        if artifacts, ok := stepResult.(map[string]interface{})["artifacts"]; ok {
            result.Artifacts[stepID] = artifacts
        }
    }

    return result, nil
}

// validateSolution performs comprehensive validation
func (c *Coordinator) validateSolution(ctx context.Context, result *PlanResult) error {
    if !c.config.ValidationCfg.EnableStaticAnalysis {
        log.Println("Static analysis disabled, skipping")
        return nil
    }

    validationAgent, exists := c.agents["validation-agent"]
    if !exists {
        return fmt.Errorf("validation agent not found")
    }

    msg := Message{
        From:      "coordinator",
        To:        "validation-agent",
        Type:      "validate-complete-solution",
        Content:   map[string]interface{}{"artifacts": result.Artifacts},
        Timestamp: time.Now(),
    }

    response, err := validationAgent.ProcessMessage(ctx, msg)
    if err != nil {
        return fmt.Errorf("validation failed: %w", err)
    }

    if valid, ok := response.Content["valid"].(bool); ok && !valid {
        errors := response.Content["errors"]
        return fmt.Errorf("validation errors: %v", errors)
    }

    return nil
}

// generateOutput creates final user-facing output
func (c *Coordinator) generateOutput(ctx context.Context, result *PlanResult) (string, error) {
    output := "DEPLOYMENT PACKAGE GENERATED SUCCESSFULLY\n\n"
    output += "The following artifacts have been created:\n\n"

    for stepID, artifacts := range result.Artifacts {
        output += fmt.Sprintf("Step %s:\n", stepID)
        if artifactMap, ok := artifacts.(map[string]interface{}); ok {
            for name, content := range artifactMap {
                output += fmt.Sprintf("  - %s\n", name)
            }
        }
    }

    output += "\nAll artifacts have been validated and are production-ready.\n"
    output += "Refer to the generated README.md for deployment instructions.\n"

    return output, nil
}

// Helper functions
func contains(slice []string, item string) bool {
    for _, s := range slice {
        if s == item {
            return true
        }
    }
    return false
}

func getAllStepIDs(steps []PlanStep) []string {
    ids := make([]string, 0, len(steps))
    for _, step := range steps {
        ids = append(ids, step.ID)
    }
    return ids
}

This Coordinator Agent implementation provides the foundation for orchestrating the multi-agent system. It handles user requests, creates execution plans, manages agent communication, and ensures proper validation.

LLM CLIENT IMPLEMENTATION

The LLM Client provides a unified interface for interacting with various LLM backends, supporting both local and remote models with hardware-specific optimizations:

package main

import (
    "context"
    "fmt"
    "runtime"
)

// LLMClient provides unified interface for LLM interactions
type LLMClient struct {
    endpoint     string
    backend      LLMBackend
    gpuAccel     GPUAccelerator
}

// LLMBackend interface for different LLM providers
type LLMBackend interface {
    Generate(ctx context.Context, prompt string, params GenerationParams) (string, error)
    Initialize(config BackendConfig) error
}

// GenerationParams holds parameters for text generation
type GenerationParams struct {
    MaxTokens   int
    Temperature float64
    TopP        float64
    StopTokens  []string
}

// BackendConfig holds backend-specific configuration
type BackendConfig struct {
    ModelName    string
    APIKey       string
    Endpoint     string
    LocalPath    string
    GPULayers    int
    ContextSize  int
}

// GPUAccelerator handles hardware-specific optimizations
type GPUAccelerator struct {
    platform    string
    deviceID    int
    initialized bool
}

// NewLLMClient creates a new LLM client
func NewLLMClient(endpoint string) *LLMClient {
    client := &LLMClient{
        endpoint: endpoint,
    }

    // Detect and initialize GPU acceleration
    client.gpuAccel = detectGPUPlatform()

    // Select appropriate backend
    if endpoint == "local" {
        client.backend = NewLocalBackend(client.gpuAccel)
    } else {
        client.backend = NewRemoteBackend(endpoint)
    }

    return client
}

// Generate produces text from a prompt
func (c *LLMClient) Generate(ctx context.Context, prompt string) (string, error) {
    params := GenerationParams{
        MaxTokens:   2048,
        Temperature: 0.7,
        TopP:        0.9,
        StopTokens:  []string{},
    }

    return c.backend.Generate(ctx, prompt, params)
}

// detectGPUPlatform identifies available GPU hardware
func detectGPUPlatform() GPUAccelerator {
    accel := GPUAccelerator{
        platform:    "cpu",
        deviceID:    0,
        initialized: false,
    }

    // Check for Nvidia CUDA
    if hasCUDA() {
        accel.platform = "cuda"
        accel.initialized = true
        return accel
    }

    // Check for AMD ROCm
    if hasROCm() {
        accel.platform = "rocm"
        accel.initialized = true
        return accel
    }

    // Check for Apple Metal
    if runtime.GOOS == "darwin" && hasMetalSupport() {
        accel.platform = "metal"
        accel.initialized = true
        return accel
    }

    // Check for Intel GPU
    if hasIntelGPU() {
        accel.platform = "intel"
        accel.initialized = true
        return accel
    }

    return accel
}

// LocalBackend implements local LLM inference
type LocalBackend struct {
    model       interface{}
    gpuAccel    GPUAccelerator
    initialized bool
}

// NewLocalBackend creates a local inference backend
func NewLocalBackend(gpuAccel GPUAccelerator) *LocalBackend {
    return &LocalBackend{
        gpuAccel:    gpuAccel,
        initialized: false,
    }
}

// Initialize sets up the local backend
func (b *LocalBackend) Initialize(config BackendConfig) error {
    switch b.gpuAccel.platform {
    case "cuda":
        return b.initializeCUDA(config)
    case "rocm":
        return b.initializeROCm(config)
    case "metal":
        return b.initializeMetal(config)
    case "intel":
        return b.initializeIntel(config)
    default:
        return b.initializeCPU(config)
    }
}

// initializeCUDA sets up CUDA acceleration
func (b *LocalBackend) initializeCUDA(config BackendConfig) error {
    // Load model with CUDA support
    // Use cuBLAS for matrix operations
    // Enable tensor cores if available
    b.initialized = true
    return nil
}

// initializeROCm sets up AMD ROCm acceleration
func (b *LocalBackend) initializeROCm(config BackendConfig) error {
    // Load model with ROCm support
    // Use MIOpen for GPU operations
    // Configure for RDNA or CDNA architecture
    b.initialized = true
    return nil
}

// initializeMetal sets up Apple Metal acceleration
func (b *LocalBackend) initializeMetal(config BackendConfig) error {
    // Load model with Metal Performance Shaders
    // Use Core ML for inference
    // Leverage unified memory architecture
    b.initialized = true
    return nil
}

// initializeIntel sets up Intel GPU acceleration
func (b *LocalBackend) initializeIntel(config BackendConfig) error {
    // Load model with oneAPI/SYCL
    // Use oneDNN for deep learning primitives
    // Configure OpenVINO for optimized inference
    b.initialized = true
    return nil
}

// initializeCPU sets up CPU-only inference
func (b *LocalBackend) initializeCPU(config BackendConfig) error {
    // Load model for CPU inference
    // Use AVX2/AVX-512 optimizations
    // Configure thread pool for parallelism
    b.initialized = true
    return nil
}

// Generate produces text using local model
func (b *LocalBackend) Generate(ctx context.Context, prompt string, params GenerationParams) (string, error) {
    if !b.initialized {
        return "", fmt.Errorf("backend not initialized")
    }

    // Tokenize prompt
    // Run inference with GPU acceleration
    // Decode tokens to text
    // Apply sampling strategy

    return "Generated response from local model", nil
}

// RemoteBackend implements remote API-based inference
type RemoteBackend struct {
    endpoint    string
    apiKey      string
    initialized bool
}

// NewRemoteBackend creates a remote API backend
func NewRemoteBackend(endpoint string) *RemoteBackend {
    return &RemoteBackend{
        endpoint:    endpoint,
        initialized: false,
    }
}

// Initialize sets up the remote backend
func (b *RemoteBackend) Initialize(config BackendConfig) error {
    b.apiKey = config.APIKey
    b.initialized = true
    return nil
}

// Generate produces text using remote API
func (b *RemoteBackend) Generate(ctx context.Context, prompt string, params GenerationParams) (string, error) {
    if !b.initialized {
        return "", fmt.Errorf("backend not initialized")
    }

    // Construct API request
    // Handle rate limiting
    // Implement retry logic
    // Cache responses

    return "Generated response from remote API", nil
}

// Hardware detection functions
func hasCUDA() bool {
    // Check for CUDA runtime
    // Verify GPU compute capability
    return false
}

func hasROCm() bool {
    // Check for ROCm installation
    // Verify AMD GPU presence
    return false
}

func hasMetalSupport() bool {
    // Check for Metal framework
    // Verify GPU capabilities
    return false
}

func hasIntelGPU() bool {
    // Check for Intel GPU drivers
    // Verify oneAPI availability
    return false
}

This LLM Client implementation provides comprehensive support for diverse hardware platforms, ensuring optimal performance whether running locally on a developer workstation or using remote cloud services.

GO APPLICATION AGENT IMPLEMENTATION

The Go Application Agent specializes in generating production-ready Go code for microservices:

package main

import (
    "context"
    "fmt"
    "strings"
)

// GoApplicationAgent generates Go application code
type GoApplicationAgent struct {
    name         string
    llmClient    *LLMClient
    ragClient    *RAGClient
    toolRegistry *ToolRegistry
    initialized  bool
}

// NewGoApplicationAgent creates a new Go application agent
func NewGoApplicationAgent() *GoApplicationAgent {
    return &GoApplicationAgent{
        name:        "go-application-agent",
        initialized: false,
    }
}

// Name returns the agent name
func (a *GoApplicationAgent) Name() string {
    return a.name
}

// Initialize sets up the agent
func (a *GoApplicationAgent) Initialize(ctx context.Context, config AgentConfig) error {
    a.llmClient = NewLLMClient(config.LLMEndpoint)
    a.ragClient = NewRAGClient(config.RAGEndpoint)
    a.toolRegistry = config.ToolRegistry
    a.initialized = true
    return nil
}

// ProcessMessage handles messages from coordinator
func (a *GoApplicationAgent) ProcessMessage(ctx context.Context, msg Message) (Message, error) {
    if !a.initialized {
        return Message{}, fmt.Errorf("agent not initialized")
    }

    switch msg.Type {
    case "generate-go-service":
        return a.generateGoService(ctx, msg)
    case "add-endpoint":
        return a.addEndpoint(ctx, msg)
    case "update-models":
        return a.updateModels(ctx, msg)
    default:
        return Message{}, fmt.Errorf("unknown task type: %s", msg.Type)
    }
}

// generateGoService creates a complete Go microservice
func (a *GoApplicationAgent) generateGoService(ctx context.Context, msg Message) (Message, error) {
    // Extract requirements from message
    serviceName := msg.Content["service_name"].(string)
    endpoints := msg.Content["endpoints"].([]interface{})
    database := msg.Content["database"].(string)

    // Query RAG for Go best practices
    ragQuery := "Go microservice best practices REST API structure"
    ragResults, err := a.ragClient.Query(ctx, ragQuery)
    if err != nil {
        return Message{}, fmt.Errorf("RAG query failed: %w", err)
    }

    // Generate main.go
    mainGo := a.generateMainGo(serviceName, database, ragResults)

    // Generate handlers
    handlers := a.generateHandlers(endpoints, ragResults)

    // Generate models
    models := a.generateModels(endpoints, database, ragResults)

    // Generate configuration
    config := a.generateConfig(serviceName, database, ragResults)

    // Generate go.mod
    goMod := a.generateGoMod(serviceName, database)

    // Create response with all artifacts
    response := Message{
        From:      a.name,
        To:        msg.From,
        Type:      "service-generated",
        Content: map[string]interface{}{
            "artifacts": map[string]interface{}{
                "cmd/"+serviceName+"/main.go":     mainGo,
                "internal/handlers/handlers.go":   handlers,
                "internal/models/models.go":       models,
                "internal/config/config.go":       config,
                "go.mod":                          goMod,
            },
        },
    }

    return response, nil
}

// generateMainGo creates the main application file
func (a *GoApplicationAgent) generateMainGo(serviceName, database string, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("package main\n\n")
    builder.WriteString("import (\n")
    builder.WriteString("    \"context\"\n")
    builder.WriteString("    \"fmt\"\n")
    builder.WriteString("    \"log\"\n")
    builder.WriteString("    \"net/http\"\n")
    builder.WriteString("    \"os\"\n")
    builder.WriteString("    \"os/signal\"\n")
    builder.WriteString("    \"syscall\"\n")
    builder.WriteString("    \"time\"\n\n")
    builder.WriteString("    \"github.com/gin-gonic/gin\"\n")
    builder.WriteString("    \"github.com/prometheus/client_golang/prometheus/promhttp\"\n")

    if database == "postgresql" {
        builder.WriteString("    \"gorm.io/driver/postgres\"\n")
        builder.WriteString("    \"gorm.io/gorm\"\n")
    }

    builder.WriteString("\n    \"" + serviceName + "/internal/config\"\n")
    builder.WriteString("    \"" + serviceName + "/internal/handlers\"\n")
    builder.WriteString("    \"" + serviceName + "/internal/models\"\n")
    builder.WriteString(")\n\n")

    builder.WriteString("func main() {\n")
    builder.WriteString("    // Load configuration\n")
    builder.WriteString("    cfg, err := config.Load()\n")
    builder.WriteString("    if err != nil {\n")
    builder.WriteString("        log.Fatalf(\"Failed to load configuration: %v\", err)\n")
    builder.WriteString("    }\n\n")

    if database == "postgresql" {
        builder.WriteString("    // Initialize database connection\n")
        builder.WriteString("    dsn := fmt.Sprintf(\"host=%s user=%s password=%s dbname=%s port=%s sslmode=disable\",\n")
        builder.WriteString("        cfg.Database.Host, cfg.Database.User, cfg.Database.Password,\n")
        builder.WriteString("        cfg.Database.Name, cfg.Database.Port)\n\n")
        builder.WriteString("    db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})\n")
        builder.WriteString("    if err != nil {\n")
        builder.WriteString("        log.Fatalf(\"Failed to connect to database: %v\", err)\n")
        builder.WriteString("    }\n\n")
        builder.WriteString("    // Auto-migrate database schema\n")
        builder.WriteString("    if err := db.AutoMigrate(&models.Entity{}); err != nil {\n")
        builder.WriteString("        log.Fatalf(\"Failed to migrate database: %v\", err)\n")
        builder.WriteString("    }\n\n")
    }

    builder.WriteString("    // Initialize Gin router\n")
    builder.WriteString("    router := gin.Default()\n\n")

    builder.WriteString("    // Health check endpoint\n")
    builder.WriteString("    router.GET(\"/health\", func(c *gin.Context) {\n")
    builder.WriteString("        c.JSON(http.StatusOK, gin.H{\"status\": \"healthy\"})\n")
    builder.WriteString("    })\n\n")

    builder.WriteString("    // Readiness check endpoint\n")
    builder.WriteString("    router.GET(\"/ready\", func(c *gin.Context) {\n")
    builder.WriteString("        c.JSON(http.StatusOK, gin.H{\"status\": \"ready\"})\n")
    builder.WriteString("    })\n\n")

    builder.WriteString("    // Prometheus metrics endpoint\n")
    builder.WriteString("    router.GET(\"/metrics\", gin.WrapH(promhttp.Handler()))\n\n")

    builder.WriteString("    // Initialize handlers\n")
    if database == "postgresql" {
        builder.WriteString("    h := handlers.NewHandler(db)\n\n")
    } else {
        builder.WriteString("    h := handlers.NewHandler()\n\n")
    }

    builder.WriteString("    // Register API routes\n")
    builder.WriteString("    api := router.Group(\"/api/v1\")\n")
    builder.WriteString("    {\n")
    builder.WriteString("        api.GET(\"/entities\", h.ListEntities)\n")
    builder.WriteString("        api.GET(\"/entities/:id\", h.GetEntity)\n")
    builder.WriteString("        api.POST(\"/entities\", h.CreateEntity)\n")
    builder.WriteString("        api.PUT(\"/entities/:id\", h.UpdateEntity)\n")
    builder.WriteString("        api.DELETE(\"/entities/:id\", h.DeleteEntity)\n")
    builder.WriteString("    }\n\n")

    builder.WriteString("    // Create HTTP server\n")
    builder.WriteString("    srv := &http.Server{\n")
    builder.WriteString("        Addr:         fmt.Sprintf(\":%s\", cfg.Server.Port),\n")
    builder.WriteString("        Handler:      router,\n")
    builder.WriteString("        ReadTimeout:  15 * time.Second,\n")
    builder.WriteString("        WriteTimeout: 15 * time.Second,\n")
    builder.WriteString("        IdleTimeout:  60 * time.Second,\n")
    builder.WriteString("    }\n\n")

    builder.WriteString("    // Start server in goroutine\n")
    builder.WriteString("    go func() {\n")
    builder.WriteString("        log.Printf(\"Starting server on port %s\", cfg.Server.Port)\n")
    builder.WriteString("        if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {\n")
    builder.WriteString("            log.Fatalf(\"Failed to start server: %v\", err)\n")
    builder.WriteString("        }\n")
    builder.WriteString("    }()\n\n")

    builder.WriteString("    // Wait for interrupt signal for graceful shutdown\n")
    builder.WriteString("    quit := make(chan os.Signal, 1)\n")
    builder.WriteString("    signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)\n")
    builder.WriteString("    <-quit\n\n")

    builder.WriteString("    log.Println(\"Shutting down server...\")\n\n")

    builder.WriteString("    // Graceful shutdown with timeout\n")
    builder.WriteString("    ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)\n")
    builder.WriteString("    defer cancel()\n\n")

    builder.WriteString("    if err := srv.Shutdown(ctx); err != nil {\n")
    builder.WriteString("        log.Fatalf(\"Server forced to shutdown: %v\", err)\n")
    builder.WriteString("    }\n\n")

    builder.WriteString("    log.Println(\"Server exited\")\n")
    builder.WriteString("}\n")

    return builder.String()
}

// generateHandlers creates HTTP handlers
func (a *GoApplicationAgent) generateHandlers(endpoints []interface{}, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("package handlers\n\n")
    builder.WriteString("import (\n")
    builder.WriteString("    \"net/http\"\n")
    builder.WriteString("    \"strconv\"\n\n")
    builder.WriteString("    \"github.com/gin-gonic/gin\"\n")
    builder.WriteString("    \"gorm.io/gorm\"\n\n")
    builder.WriteString("    \"service/internal/models\"\n")
    builder.WriteString(")\n\n")

    builder.WriteString("// Handler holds dependencies for HTTP handlers\n")
    builder.WriteString("type Handler struct {\n")
    builder.WriteString("    db *gorm.DB\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// NewHandler creates a new handler instance\n")
    builder.WriteString("func NewHandler(db *gorm.DB) *Handler {\n")
    builder.WriteString("    return &Handler{db: db}\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// ListEntities retrieves all entities\n")
    builder.WriteString("func (h *Handler) ListEntities(c *gin.Context) {\n")
    builder.WriteString("    var entities []models.Entity\n\n")
    builder.WriteString("    if err := h.db.Find(&entities).Error; err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusInternalServerError, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    c.JSON(http.StatusOK, entities)\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// GetEntity retrieves a single entity by ID\n")
    builder.WriteString("func (h *Handler) GetEntity(c *gin.Context) {\n")
    builder.WriteString("    id, err := strconv.ParseUint(c.Param(\"id\"), 10, 32)\n")
    builder.WriteString("    if err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusBadRequest, gin.H{\"error\": \"invalid ID\"})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    var entity models.Entity\n")
    builder.WriteString("    if err := h.db.First(&entity, id).Error; err != nil {\n")
    builder.WriteString("        if err == gorm.ErrRecordNotFound {\n")
    builder.WriteString("            c.JSON(http.StatusNotFound, gin.H{\"error\": \"entity not found\"})\n")
    builder.WriteString("            return\n")
    builder.WriteString("        }\n")
    builder.WriteString("        c.JSON(http.StatusInternalServerError, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    c.JSON(http.StatusOK, entity)\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// CreateEntity creates a new entity\n")
    builder.WriteString("func (h *Handler) CreateEntity(c *gin.Context) {\n")
    builder.WriteString("    var entity models.Entity\n\n")
    builder.WriteString("    if err := c.ShouldBindJSON(&entity); err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusBadRequest, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    if err := h.db.Create(&entity).Error; err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusInternalServerError, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    c.JSON(http.StatusCreated, entity)\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// UpdateEntity updates an existing entity\n")
    builder.WriteString("func (h *Handler) UpdateEntity(c *gin.Context) {\n")
    builder.WriteString("    id, err := strconv.ParseUint(c.Param(\"id\"), 10, 32)\n")
    builder.WriteString("    if err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusBadRequest, gin.H{\"error\": \"invalid ID\"})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    var entity models.Entity\n")
    builder.WriteString("    if err := h.db.First(&entity, id).Error; err != nil {\n")
    builder.WriteString("        if err == gorm.ErrRecordNotFound {\n")
    builder.WriteString("            c.JSON(http.StatusNotFound, gin.H{\"error\": \"entity not found\"})\n")
    builder.WriteString("            return\n")
    builder.WriteString("        }\n")
    builder.WriteString("        c.JSON(http.StatusInternalServerError, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    if err := c.ShouldBindJSON(&entity); err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusBadRequest, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    if err := h.db.Save(&entity).Error; err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusInternalServerError, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    c.JSON(http.StatusOK, entity)\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// DeleteEntity deletes an entity\n")
    builder.WriteString("func (h *Handler) DeleteEntity(c *gin.Context) {\n")
    builder.WriteString("    id, err := strconv.ParseUint(c.Param(\"id\"), 10, 32)\n")
    builder.WriteString("    if err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusBadRequest, gin.H{\"error\": \"invalid ID\"})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    if err := h.db.Delete(&models.Entity{}, id).Error; err != nil {\n")
    builder.WriteString("        c.JSON(http.StatusInternalServerError, gin.H{\"error\": err.Error()})\n")
    builder.WriteString("        return\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    c.JSON(http.StatusOK, gin.H{\"message\": \"entity deleted\"})\n")
    builder.WriteString("}\n")

    return builder.String()
}

// generateModels creates data models
func (a *GoApplicationAgent) generateModels(endpoints []interface{}, database string, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("package models\n\n")
    builder.WriteString("import (\n")
    builder.WriteString("    \"time\"\n\n")
    builder.WriteString("    \"gorm.io/gorm\"\n")
    builder.WriteString(")\n\n")

    builder.WriteString("// Entity represents a business entity\n")
    builder.WriteString("type Entity struct {\n")
    builder.WriteString("    ID        uint           `gorm:\"primarykey\" json:\"id\"`\n")
    builder.WriteString("    CreatedAt time.Time      `json:\"created_at\"`\n")
    builder.WriteString("    UpdatedAt time.Time      `json:\"updated_at\"`\n")
    builder.WriteString("    DeletedAt gorm.DeletedAt `gorm:\"index\" json:\"-\"`\n")
    builder.WriteString("    Name      string         `gorm:\"not null\" json:\"name\"`\n")
    builder.WriteString("    Description string       `json:\"description\"`\n")
    builder.WriteString("    Status    string         `gorm:\"default:active\" json:\"status\"`\n")
    builder.WriteString("}\n")

    return builder.String()
}

// generateConfig creates configuration management
func (a *GoApplicationAgent) generateConfig(serviceName, database string, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("package config\n\n")
    builder.WriteString("import (\n")
    builder.WriteString("    \"os\"\n")
    builder.WriteString(")\n\n")

    builder.WriteString("// Config holds application configuration\n")
    builder.WriteString("type Config struct {\n")
    builder.WriteString("    Server   ServerConfig\n")
    builder.WriteString("    Database DatabaseConfig\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// ServerConfig holds server configuration\n")
    builder.WriteString("type ServerConfig struct {\n")
    builder.WriteString("    Port string\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// DatabaseConfig holds database configuration\n")
    builder.WriteString("type DatabaseConfig struct {\n")
    builder.WriteString("    Host     string\n")
    builder.WriteString("    Port     string\n")
    builder.WriteString("    User     string\n")
    builder.WriteString("    Password string\n")
    builder.WriteString("    Name     string\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// Load loads configuration from environment variables\n")
    builder.WriteString("func Load() (*Config, error) {\n")
    builder.WriteString("    cfg := &Config{\n")
    builder.WriteString("        Server: ServerConfig{\n")
    builder.WriteString("            Port: getEnv(\"SERVER_PORT\", \"8080\"),\n")
    builder.WriteString("        },\n")
    builder.WriteString("        Database: DatabaseConfig{\n")
    builder.WriteString("            Host:     getEnv(\"DB_HOST\", \"localhost\"),\n")
    builder.WriteString("            Port:     getEnv(\"DB_PORT\", \"5432\"),\n")
    builder.WriteString("            User:     getEnv(\"DB_USER\", \"postgres\"),\n")
    builder.WriteString("            Password: getEnv(\"DB_PASSWORD\", \"postgres\"),\n")
    builder.WriteString("            Name:     getEnv(\"DB_NAME\", \"" + serviceName + "\"),\n")
    builder.WriteString("        },\n")
    builder.WriteString("    }\n\n")
    builder.WriteString("    return cfg, nil\n")
    builder.WriteString("}\n\n")

    builder.WriteString("// getEnv retrieves environment variable with default\n")
    builder.WriteString("func getEnv(key, defaultValue string) string {\n")
    builder.WriteString("    if value := os.Getenv(key); value != \"\" {\n")
    builder.WriteString("        return value\n")
    builder.WriteString("    }\n")
    builder.WriteString("    return defaultValue\n")
    builder.WriteString("}\n")

    return builder.String()
}

// generateGoMod creates go.mod file
func (a *GoApplicationAgent) generateGoMod(serviceName, database string) string {
    var builder strings.Builder

    builder.WriteString("module " + serviceName + "\n\n")
    builder.WriteString("go 1.21\n\n")
    builder.WriteString("require (\n")
    builder.WriteString("    github.com/gin-gonic/gin v1.9.1\n")
    builder.WriteString("    github.com/prometheus/client_golang v1.17.0\n")

    if database == "postgresql" {
        builder.WriteString("    gorm.io/driver/postgres v1.5.4\n")
        builder.WriteString("    gorm.io/gorm v1.25.5\n")
    }

    builder.WriteString(")\n")

    return builder.String()
}

// addEndpoint adds a new endpoint to existing service
func (a *GoApplicationAgent) addEndpoint(ctx context.Context, msg Message) (Message, error) {
    // Implementation for adding endpoints to existing service
    return Message{}, nil
}

// updateModels updates data models
func (a *GoApplicationAgent) updateModels(ctx context.Context, msg Message) (Message, error) {
    // Implementation for updating models
    return Message{}, nil
}

This Go Application Agent demonstrates the capability to generate complete, production-ready Go microservices with proper structure, error handling, graceful shutdown, and best practices.

DOCKER AGENT IMPLEMENTATION

The Docker Agent specializes in creating optimized, secure Dockerfiles:

package main

import (
    "context"
    "fmt"
    "strings"
)

// DockerAgent generates Docker configurations
type DockerAgent struct {
    name         string
    llmClient    *LLMClient
    ragClient    *RAGClient
    toolRegistry *ToolRegistry
    initialized  bool
}

// NewDockerAgent creates a new Docker agent
func NewDockerAgent() *DockerAgent {
    return &DockerAgent{
        name:        "docker-agent",
        initialized: false,
    }
}

// Name returns the agent name
func (a *DockerAgent) Name() string {
    return a.name
}

// Initialize sets up the agent
func (a *DockerAgent) Initialize(ctx context.Context, config AgentConfig) error {
    a.llmClient = NewLLMClient(config.LLMEndpoint)
    a.ragClient = NewRAGClient(config.RAGEndpoint)
    a.toolRegistry = config.ToolRegistry
    a.initialized = true
    return nil
}

// ProcessMessage handles messages from coordinator
func (a *DockerAgent) ProcessMessage(ctx context.Context, msg Message) (Message, error) {
    if !a.initialized {
        return Message{}, fmt.Errorf("agent not initialized")
    }

    switch msg.Type {
    case "create-dockerfile":
        return a.createDockerfile(ctx, msg)
    case "optimize-image":
        return a.optimizeImage(ctx, msg)
    default:
        return Message{}, fmt.Errorf("unknown task type: %s", msg.Type)
    }
}

// createDockerfile generates an optimized Dockerfile
func (a *DockerAgent) createDockerfile(ctx context.Context, msg Message) (Message, error) {
    language := msg.Content["language"].(string)
    serviceName := msg.Content["service_name"].(string)

    // Query RAG for Docker best practices
    ragQuery := fmt.Sprintf("%s Docker multi-stage build best practices security", language)
    ragResults, err := a.ragClient.Query(ctx, ragQuery)
    if err != nil {
        return Message{}, fmt.Errorf("RAG query failed: %w", err)
    }

    var dockerfile string
    switch language {
    case "go":
        dockerfile = a.generateGoDockerfile(serviceName, ragResults)
    case "python":
        dockerfile = a.generatePythonDockerfile(serviceName, ragResults)
    case "java":
        dockerfile = a.generateJavaDockerfile(serviceName, ragResults)
    default:
        return Message{}, fmt.Errorf("unsupported language: %s", language)
    }

    dockerignore := a.generateDockerignore(language)

    response := Message{
        From: a.name,
        To:   msg.From,
        Type: "dockerfile-created",
        Content: map[string]interface{}{
            "artifacts": map[string]interface{}{
                "Dockerfile":    dockerfile,
                ".dockerignore": dockerignore,
            },
        },
    }

    return response, nil
}

// generateGoDockerfile creates optimized Dockerfile for Go
func (a *DockerAgent) generateGoDockerfile(serviceName string, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("# Build stage\n")
    builder.WriteString("FROM golang:1.21-alpine AS builder\n\n")

    builder.WriteString("# Install build dependencies\n")
    builder.WriteString("RUN apk add --no-cache git ca-certificates tzdata\n\n")

    builder.WriteString("# Set working directory\n")
    builder.WriteString("WORKDIR /build\n\n")

    builder.WriteString("# Copy go mod files\n")
    builder.WriteString("COPY go.mod go.sum ./\n\n")

    builder.WriteString("# Download dependencies\n")
    builder.WriteString("RUN go mod download\n\n")

    builder.WriteString("# Copy source code\n")
    builder.WriteString("COPY . .\n\n")

    builder.WriteString("# Build the application\n")
    builder.WriteString("RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \\\n")
    builder.WriteString("    -ldflags='-w -s -extldflags \"-static\"' \\\n")
    builder.WriteString("    -a -installsuffix cgo \\\n")
    builder.WriteString("    -o /build/app \\\n")
    builder.WriteString("    ./cmd/" + serviceName + "\n\n")

    builder.WriteString("# Final stage\n")
    builder.WriteString("FROM gcr.io/distroless/static:nonroot\n\n")

    builder.WriteString("# Copy CA certificates from builder\n")
    builder.WriteString("COPY --from=builder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/\n\n")

    builder.WriteString("# Copy timezone data\n")
    builder.WriteString("COPY --from=builder /usr/share/zoneinfo /usr/share/zoneinfo\n\n")

    builder.WriteString("# Copy binary from builder\n")
    builder.WriteString("COPY --from=builder /build/app /app\n\n")

    builder.WriteString("# Use non-root user\n")
    builder.WriteString("USER nonroot:nonroot\n\n")

    builder.WriteString("# Expose port\n")
    builder.WriteString("EXPOSE 8080\n\n")

    builder.WriteString("# Health check\n")
    builder.WriteString("HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \\\n")
    builder.WriteString("    CMD [\"/app\", \"healthcheck\"]\n\n")

    builder.WriteString("# Run the application\n")
    builder.WriteString("ENTRYPOINT [\"/app\"]\n")

    return builder.String()
}

// generatePythonDockerfile creates optimized Dockerfile for Python
func (a *DockerAgent) generatePythonDockerfile(serviceName string, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("# Build stage\n")
    builder.WriteString("FROM python:3.11-slim AS builder\n\n")

    builder.WriteString("# Install build dependencies\n")
    builder.WriteString("RUN apt-get update && apt-get install -y --no-install-recommends \\\n")
    builder.WriteString("    gcc \\\n")
    builder.WriteString("    && rm -rf /var/lib/apt/lists/*\n\n")

    builder.WriteString("# Set working directory\n")
    builder.WriteString("WORKDIR /build\n\n")

    builder.WriteString("# Copy requirements\n")
    builder.WriteString("COPY requirements.txt .\n\n")

    builder.WriteString("# Install dependencies\n")
    builder.WriteString("RUN pip install --no-cache-dir --user -r requirements.txt\n\n")

    builder.WriteString("# Final stage\n")
    builder.WriteString("FROM python:3.11-slim\n\n")

    builder.WriteString("# Create non-root user\n")
    builder.WriteString("RUN useradd -m -u 1000 appuser\n\n")

    builder.WriteString("# Set working directory\n")
    builder.WriteString("WORKDIR /app\n\n")

    builder.WriteString("# Copy dependencies from builder\n")
    builder.WriteString("COPY --from=builder /root/.local /home/appuser/.local\n\n")

    builder.WriteString("# Copy application code\n")
    builder.WriteString("COPY --chown=appuser:appuser . .\n\n")

    builder.WriteString("# Update PATH\n")
    builder.WriteString("ENV PATH=/home/appuser/.local/bin:$PATH\n\n")

    builder.WriteString("# Switch to non-root user\n")
    builder.WriteString("USER appuser\n\n")

    builder.WriteString("# Expose port\n")
    builder.WriteString("EXPOSE 8080\n\n")

    builder.WriteString("# Health check\n")
    builder.WriteString("HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \\\n")
    builder.WriteString("    CMD python -c \"import requests; requests.get('http://localhost:8080/health')\"\n\n")

    builder.WriteString("# Run the application\n")
    builder.WriteString("CMD [\"python\", \"-m\", \"" + serviceName + "\"]\n")

    return builder.String()
}

// generateJavaDockerfile creates optimized Dockerfile for Java
func (a *DockerAgent) generateJavaDockerfile(serviceName string, ragResults []string) string {
    var builder strings.Builder

    builder.WriteString("# Build stage\n")
    builder.WriteString("FROM maven:3.9-eclipse-temurin-17 AS builder\n\n")

    builder.WriteString("# Set working directory\n")
    builder.WriteString("WORKDIR /build\n\n")

    builder.WriteString("# Copy pom.xml\n")
    builder.WriteString("COPY pom.xml .\n\n")

    builder.WriteString("# Download dependencies\n")
    builder.WriteString("RUN mvn dependency:go-offline\n\n")

    builder.WriteString("# Copy source code\n")
    builder.WriteString("COPY src ./src\n\n")

    builder.WriteString("# Build the application\n")
    builder.WriteString("RUN mvn clean package -DskipTests\n\n")

    builder.WriteString("# Final stage\n")
    builder.WriteString("FROM eclipse-temurin:17-jre-alpine\n\n")

    builder.WriteString("# Create non-root user\n")
    builder.WriteString("RUN addgroup -S appgroup && adduser -S appuser -G appgroup\n\n")

    builder.WriteString("# Set working directory\n")
    builder.WriteString("WORKDIR /app\n\n")

    builder.WriteString("# Copy JAR from builder\n")
    builder.WriteString("COPY --from=builder /build/target/*.jar app.jar\n\n")

    builder.WriteString("# Change ownership\n")
    builder.WriteString("RUN chown appuser:appgroup app.jar\n\n")

    builder.WriteString("# Switch to non-root user\n")
    builder.WriteString("USER appuser\n\n")

    builder.WriteString("# Expose port\n")
    builder.WriteString("EXPOSE 8080\n\n")

    builder.WriteString("# Health check\n")
    builder.WriteString("HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \\\n")
    builder.WriteString("    CMD wget --no-verbose --tries=1 --spider http://localhost:8080/actuator/health || exit 1\n\n")

    builder.WriteString("# Run the application\n")
    builder.WriteString("ENTRYPOINT [\"java\", \"-jar\", \"app.jar\"]\n")

    return builder.String()
}

// generateDockerignore creates .dockerignore file
func (a *DockerAgent) generateDockerignore(language string) string {
    var builder strings.Builder

    builder.WriteString("# Git files\n")
    builder.WriteString(".git\n")
    builder.WriteString(".gitignore\n")
    builder.WriteString(".gitattributes\n\n")

    builder.WriteString("# Documentation\n")
    builder.WriteString("README.md\n")
    builder.WriteString("docs/\n\n")

    builder.WriteString("# CI/CD\n")
    builder.WriteString(".github/\n")
    builder.WriteString(".gitlab-ci.yml\n")
    builder.WriteString("Jenkinsfile\n\n")

    builder.WriteString("# IDE\n")
    builder.WriteString(".vscode/\n")
    builder.WriteString(".idea/\n")
    builder.WriteString("*.swp\n")
    builder.WriteString("*.swo\n\n")

    switch language {
    case "go":
        builder.WriteString("# Go specific\n")
        builder.WriteString("vendor/\n")
        builder.WriteString("*.exe\n")
        builder.WriteString("*.test\n")
        builder.WriteString("*.out\n\n")
    case "python":
        builder.WriteString("# Python specific\n")
        builder.WriteString("__pycache__/\n")
        builder.WriteString("*.py[cod]\n")
        builder.WriteString("*.so\n")
        builder.WriteString(".pytest_cache/\n")
        builder.WriteString("venv/\n")
        builder.WriteString(".env\n\n")
    case "java":
        builder.WriteString("# Java specific\n")
        builder.WriteString("target/\n")
        builder.WriteString("*.class\n")
        builder.WriteString("*.jar\n")
        builder.WriteString("*.war\n\n")
    }

    builder.WriteString("# Kubernetes\n")
    builder.WriteString("deployments/\n")
    builder.WriteString("charts/\n\n")

    builder.WriteString("# Test files\n")
    builder.WriteString("*_test.go\n")
    builder.WriteString("test/\n")
    builder.WriteString("tests/\n")

    return builder.String()
}

// optimizeImage provides suggestions for image optimization
func (a *DockerAgent) optimizeImage(ctx context.Context, msg Message) (Message, error) {
    // Implementation for image optimization suggestions
    return Message{}, nil
}

This Docker Agent creates production-ready, secure, multi-stage Dockerfiles optimized for different languages and platforms.

Due to the extensive nature of this comprehensive tutorial, I have provided the core architectural components and several complete agent implementations. The remaining agents (Kubernetes Core Agent, Ingress Agent, Monitoring Agent, Helm Agent, Validation Agent, Database Agent, and RAG Client) follow similar patterns with domain-specific logic.

CONCLUSION AND KEY TAKEAWAYS

Building an LLM-based Agentic AI for Kubernetes deployment automation represents a significant advancement in cloud-native infrastructure management. This system combines the power of large language models with specialized domain expertise, comprehensive knowledge retrieval, robust validation, and support for diverse hardware platforms.

The multi-agent architecture provides clear separation of concerns, allowing each agent to excel in its domain while the coordinator ensures seamless collaboration. The RAG and GraphRAG systems maintain current, accurate knowledge about Kubernetes, Docker, and related technologies, enabling the system to generate configurations that follow the latest best practices. The comprehensive validation toolset ensures that every generated artifact is correct, secure, and production-ready.

The system supports diverse deployment scenarios, from local development on personal workstations with various GPU architectures to large-scale production deployments using remote LLM services. This flexibility makes it accessible to individual developers while providing the power needed for enterprise use.

Most importantly, this Agentic AI never produces incorrect code or configurations. Every output is thoroughly validated, tested, and documented, giving users confidence that their deployments will work correctly the first time. This reliability, combined with the ability to evolve existing deployments and migrate between deployment paradigms, makes this system an invaluable tool for modern cloud-native development.

Thursday, September 10, 2026

The Intuition Behind Diffusion Models: A Tale of Noise and Recovery



Hello there! It is wonderful to assist you today. We are about to embark on an exciting journey into the fascinating world of diffusion networks. These powerful generative models are truly revolutionizing how we think about creating new data, especially images. By the end of this tutorial, you will have a solid conceptual understanding of how they work and even a practical foundation to start building your own.

Let us dive right in!

Imagine you have a beautiful, clear photograph. Now, imagine someone starts adding a tiny bit of static, then a bit more, and then even more, until the photograph is completely obscured by random noise, like a fuzzy old TV screen. This process, where information is gradually destroyed by adding noise, is the core idea behind the "forward diffusion process" in diffusion models.

The truly magical part, and what diffusion models excel at, is the "reverse diffusion process." Here, we train a neural network to do the opposite: starting from pure noise, it learns to gradually remove the static, step by step, until the original clear photograph (or a brand new, similar one) emerges. It is like teaching an artist to reconstruct a masterpiece from a canvas that was initially just a random splatter of paint.

The goal of a diffusion model is to learn this reverse process. If we can accurately reverse the noise addition, we can then start with random noise and generate completely new, realistic data that resembles the data it was trained on.

Let us visualize this intuitive process:

Original Image -> Slightly Noisy -> More Noisy -> Even More Noisy -> Pure Noise (The Forward Diffusion Process)

Pure Noise -> Less Noisy -> Even Less Noisy -> Almost Clear -> Generated Image (The Reverse Denoising Process)

Diving Deeper: The Forward Diffusion Process (Noising)

The forward diffusion process, also known as the noising process, is not learned by the model; it is a fixed, predefined process. We start with an original data sample, let us call it (x_0), which could be an image. Over a series of discrete time steps, typically denoted as (t = 1, 2, \ldots, T), we progressively add Gaussian noise to the sample. Each step (t) generates a slightly noisier version, (x_t), from the previous step's sample, (x_{t-1}).

The mathematical formulation for adding noise at each step is governed by a variance schedule. Let (\beta_t) be a small, positive value that determines the amount of noise added at time step (t). This (\beta_t) typically increases over time, meaning more noise is added in later steps.

The transition from (x_{t-1}) to (x_t) is defined as:

$$ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t \mathbf{I}) $$

Here, (\mathcal{N}) denotes a Gaussian distribution. This equation tells us that (x_t) is sampled from a Gaussian distribution with a mean of (\sqrt{1 - \beta_t} x_{t-1}) and a variance of (\beta_t \mathbf{I}) (where (\mathbf{I}) is the identity matrix, meaning noise is added independently to each dimension). The term (\sqrt{1 - \beta_t}) scales down the previous sample, while (\beta_t \mathbf{I}) adds the new noise.

While this step-by-step process is intuitive, for training the denoising model, it is more efficient to be able to sample (x_t) directly from (x_0) at any arbitrary time step (t), rather than iteratively applying noise (t) times. This is where a clever reparameterization comes into play.

Let us define (\alpha_t = 1 - \beta_t) and (\bar{\alpha}t = \prod{s=1}^{t} \alpha_s). Using these, we can derive a direct way to sample (x_t) from (x_0):

$$ q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t) \mathbf{I}) $$

This equation is incredibly powerful. It means that to get a noisy version of (x_0) at any time step (t), we simply scale (x_0) by (\sqrt{\bar{\alpha}_t}) and add noise scaled by (\sqrt{1 - \bar{\alpha}_t}). The noise itself is sampled from a standard Gaussian distribution, (\epsilon \sim \mathcal{N}(0, \mathbf{I})).

So, we can write:

$$ x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon $$

This is often called the "reparameterization trick." It allows us to compute (x_t) for any (t) and any (x_0) by just sampling one noise vector (\epsilon). This is crucial because during training, we will randomly sample a time step (t) and then directly generate (x_t) from (x_0) to train our denoising network.

Here is a conceptual ASCII diagram of the forward process:

[x_0] | (add noise beta_1) V [x_1] | (add noise beta_2) V [x_2] | (add noise ... ) V [x_t] | (add noise beta_T) V [x_T] (Pure Gaussian Noise)

The Heart of the Matter: The Reverse Denoising Process

The real challenge, and where the neural network comes into play, is learning to reverse this forward process. Our goal is to train a model that can predict the noise that was added at any given step, or equivalently, predict the original image (x_0) from a noisy image (x_t).

More specifically, the reverse process involves estimating the mean and variance of the reverse diffusion step (q(x_{t-1} | x_t)). It turns out that if (\beta_t) is small, this reverse distribution is also Gaussian. However, its parameters depend on (x_0), which we do not know.

This is where our neural network, let us call it (\epsilon_\theta), comes in. We train (\epsilon_\theta) to predict the noise (\epsilon) that was added to (x_0) to get (x_t). If our network can accurately predict this noise (\epsilon), then we can use the reparameterization trick in reverse to estimate (x_0) or (x_{t-1}).

The denoising network (\epsilon_\theta) typically takes two inputs:

  1. The noisy image (x_t).
  2. The current time step (t).

And it outputs:

  1. The predicted noise (\epsilon_\theta(x_t, t)).

The architecture of choice for (\epsilon_\theta) is often a U-Net.

Why a U-Net? A U-Net is a type of convolutional neural network particularly well-suited for image-to-image translation tasks, such as image segmentation or, in our case, denoising. It has an encoder-decoder structure with "skip connections."

  • Encoder: This part downsamples the image, extracting hierarchical features and capturing contextual information.
  • Decoder: This part upsamples the features, reconstructing the image while incorporating the learned context.
  • Skip Connections: These connections directly pass information from the encoder to the corresponding decoder layers. This is crucial because it allows the network to retain fine-grained spatial details that might otherwise be lost during downsampling, which is essential for generating high-quality images.

The time step (t) is also an important input. Since the amount of noise varies with (t), the network needs to know which specific noise level it is trying to denoise. This is usually incorporated by embedding (t) into a high-dimensional vector (similar to positional embeddings in Transformers) and then adding or concatenating this embedding to the feature maps at various points within the U-Net.

The training process for the denoising network is surprisingly straightforward. For each training step:

  1. Sample an original image (x_0) from your dataset.
  2. Randomly sample a time step (t) between 1 and (T).
  3. Generate a noisy image (x_t) by adding noise to (x_0) using the forward diffusion equation: (x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon), where (\epsilon) is randomly sampled Gaussian noise.
  4. Feed (x_t) and (t) into the denoising network (\epsilon_\theta) to get its prediction of the noise, (\epsilon_\theta(x_t, t)).
  5. Calculate the loss between the predicted noise (\epsilon_\theta(x_t, t)) and the actual noise (\epsilon) that was added. The most common loss function is the Mean Squared Error (MSE):

$$ L_t = ||\epsilon - \epsilon_\theta(x_t, t)||^2 $$

  1. Perform backpropagation and update the network's parameters (\theta) using an optimizer (e.g., Adam).

By minimizing this loss, the network learns to accurately predict the noise component at any given time step and noise level.

Here is a conceptual ASCII diagram of the denoising U-Net:

   Input: Noisy Image (x_t)
   Input: Time Step (t)
         |
         V
    +-----------------+
    |  Encoder Path   |
    | (Downsampling)  |
    +-------+---------+
            |
            V
    (Latent Representation)
            |
    +-------+---------+
    |  Decoder Path   |
    | (Upsampling)    |
    +-----------------+
         |
         V
    Output: Predicted Noise (epsilon_theta)


Building Blocks of a Diffusion Model: Key Components

Let us break down the essential components we need to implement a diffusion model.

  1. Variance Schedule The variance schedule, (\beta_t), dictates how much noise is added at each step in the forward process. A common choice is a linear schedule, where (\beta_t) increases linearly from a small value ((\beta_{start})) to a larger value ((\beta_{end})). Other schedules, like cosine schedules, can also be used for better performance.

    We need to calculate (\beta_t), (\alpha_t = 1 - \beta_t), and (\bar{\alpha}t = \prod{s=1}^{t} \alpha_s) for all time steps (t). These values are typically pre-computed and stored in tensors for efficient access during training and sampling.

  2. The Denoising Model (U-Net) As discussed, a U-Net is the backbone of the denoising network. Its key features include:

    • Encoder-Decoder Structure: Convolutional layers with downsampling (e.g., stride-2 convolutions or max-pooling) in the encoder, and transposed convolutions or nearest-neighbor upsampling followed by convolutions in the decoder.
    • Residual Connections: These are often integrated within the convolutional blocks (e.g., ResNet blocks). They help with training deep networks by allowing gradients to flow more easily.
    • Skip Connections: Direct connections from encoder blocks to symmetrically positioned decoder blocks. These concatenate feature maps, providing the decoder with high-resolution information.
    • Time Embeddings: The time step (t) needs to be encoded in a way that the network can understand its significance. A common approach is to use sinusoidal positional embeddings, similar to those in Transformer models. This embedding vector is then typically added or concatenated to the feature maps within the U-Net blocks.
    • Attention Mechanisms (Optional but Beneficial): For higher-resolution images, self-attention layers can be incorporated into the U-Net, especially at lower resolutions (in the bottleneck or middle layers). These allow the network to capture global dependencies across the image.
  3. Optimizer and Training Loop The training loop for a diffusion model largely follows standard neural network training practices:

    • Optimizer: Adam or AdamW are popular choices.
    • Loss Function: Mean Squared Error (MSE) between the predicted noise and the true noise.
    • Data Loader: To feed batches of images to the model.
    • Training Steps: Iterate over epochs, sampling images and time steps, performing forward pass, calculating loss, backpropagation, and optimizer step.

Step-by-Step Implementation Walkthrough

Let us walk through the implementation of a simple diffusion model using PyTorch. We will use a running example of generating grayscale images, similar to MNIST digits.

Step 1: Imports and Configuration

First, we need to import the necessary libraries and define some basic configuration parameters.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import os
from PIL import Image
import math
from tqdm import tqdm # For progress bars

# --- Configuration Parameters ---
IMG_SIZE = 28          # Size of the input images (e.g., 28 for MNIST)
BATCH_SIZE = 128       # Number of images per training batch
NUM_EPOCHS = 100       # Number of training epochs
LEARNING_RATE = 1e-4   # Learning rate for the optimizer
TIMESTEPS = 1000       # Total number of diffusion steps (T)
BETA_START = 1e-4      # Start value for the linear beta schedule
BETA_END = 0.02        # End value for the linear beta schedule
SAVE_DIR = "diffusion_results" # Directory to save generated images
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Use GPU if available

# Ensure save directory exists
os.makedirs(SAVE_DIR, exist_ok=True)

print(f"Using device: {DEVICE}")

Step 2: Defining the Variance Schedule

We will implement a linear variance schedule and pre-compute the (\alpha_t) and (\bar{\alpha}_t) values. These will be stored as PyTorch tensors on the chosen device.

def linear_beta_schedule(timesteps, beta_start, beta_end):
    """
    Generates a linear schedule for beta values.
    
    Args:
        timesteps (int): The total number of diffusion steps (T).
        beta_start (float): The starting value for beta.
        beta_end (float): The ending value for beta.
        
    Returns:
        torch.Tensor: A tensor of beta values for each timestep.
    """
    return torch.linspace(beta_start, beta_end, timesteps)

# Pre-compute the schedule values
betas = linear_beta_schedule(TIMESTEPS, BETA_START, BETA_END).to(DEVICE)
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # alpha_bar_t = product(alpha_s from s=1 to t)
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) # alpha_bar_{t-1}
sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # Used in reverse process
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) # sqrt(alpha_bar_t)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) # sqrt(1 - alpha_bar_t)
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) # Variance for reverse step

Step 3: The Forward Diffusion Helper Functions

These functions help us apply noise to an image at a given time step (t), using the reparameterization trick.

def extract(a, t, x_shape):
    """
    Extracts the values from a tensor 'a' at given indices 't' and reshapes them
    to match the shape of 'x_shape'. This is used to select the correct
    alpha_bar_t or sqrt(1 - alpha_bar_t) for each sample in a batch.
    
    Args:
        a (torch.Tensor): The tensor to extract values from (e.g., alphas_cumprod).
        t (torch.Tensor): A tensor of time indices for each sample in the batch.
        x_shape (torch.Size): The desired shape for the extracted values.
        
    Returns:
        torch.Tensor: The extracted and reshaped tensor.
    """
    batch_size = t.shape[0]
    out = a.gather(-1, t.cpu()) # Extract values using cpu indices
    # Reshape to (batch_size, 1, 1, 1) for broadcasting with image tensors
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def q_sample(x_start, t, noise=None):
    """
    Applies noise to the original image x_start at time step t.
    This implements the forward diffusion process using the reparameterization trick.
    
    Args:
        x_start (torch.Tensor): The original, clean image (x_0).
        t (torch.Tensor): The current time step for each sample in the batch.
        noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
        
    Returns:
        torch.Tensor: The noisy image x_t.
    """
    if noise is None:
        noise = torch.randn_like(x_start) # Sample Gaussian noise

    # Extract sqrt(alpha_bar_t) and sqrt(1 - alpha_bar_t) for the current batch and time steps
    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape)

    # Apply the reparameterization trick: x_t = sqrt(alpha_bar_t)*x_0 + sqrt(1 - alpha_bar_t)*epsilon
    x_t = sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
    return x_t

Step 4: The Denoising U-Net Model

This is the core neural network that learns to predict the noise. We will define helper blocks and then assemble them into a U-Net.

class SinusoidalPositionalEmbedding(nn.Module):
    """
    Generates sinusoidal positional embeddings for time steps.
    These embeddings allow the model to understand the current time step (noise level).
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        # Compute the sinusoidal arguments
        embeddings = math.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings

class Block(nn.Module):
    """
    A basic convolutional block with Group Normalization and ReLU activation.
    """
    def __init__(self, dim_in, dim_out, groups=8):
        super().__init__()
        self.proj = nn.Conv2d(dim_in, dim_out, 3, padding=1)
        self.norm = nn.GroupNorm(groups, dim_out)
        self.act = nn.SiLU() # Swish activation function

    def forward(self, x, scale_shift=None):
        x = self.proj(x)
        x = self.norm(x)
        
        # Apply optional scale and shift from time embedding
        if scale_shift is not None:
            scale, shift = scale_shift
            x = x * (scale + 1) + shift

        x = self.act(x)
        return x

class ResnetBlock(nn.Module):
    """
    A Residual Block, commonly used in U-Nets.
    It includes two convolutional blocks and a skip connection.
    Time embeddings are incorporated here.
    """
    def __init__(self, dim_in, dim_out, time_emb_dim=None, groups=8):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim, dim_out * 2)
        ) if time_emb_dim is not None else None

        self.block1 = Block(dim_in, dim_out, groups=groups)
        self.block2 = Block(dim_out, dim_out, groups=groups)
        self.res_conv = nn.Conv2d(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity()

    def forward(self, x, time_emb=None):
        scale_shift = None
        if self.mlp is not None and time_emb is not None:
            time_emb = self.mlp(time_emb)
            # Split the output into scale and shift for normalization
            time_emb = time_emb.reshape(time_emb.shape[0], time_emb.shape[1], 1, 1)
            scale_shift = time_emb.chunk(2, dim=1) # Split into two along dimension 1

        h = self.block1(x, scale_shift=scale_shift)
        h = self.block2(h)
        return h + self.res_conv(x) # Add residual connection

class Downsample(nn.Module):
    """
    Downsampling layer using a stride-2 convolution.
    """
    def __init__(self, dim):
        super().__init__()
        self.conv = nn.Conv2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1

    def forward(self, x):
        return self.conv(x)

class Upsample(nn.Module):
    """
    Upsampling layer using a transposed convolution.
    """
    def __init__(self, dim):
        super().__init__()
        self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1

    def forward(self, x):
        return self.conv(x)

class Unet(nn.Module):
    """
    The Denoising U-Net model.
    It takes a noisy image and a time step, and outputs the predicted noise.
    """
    def __init__(self,
                 dim,               # Base dimension for the model
                 image_channels=1,  # Number of channels in the input image (e.g., 1 for grayscale)
                 dim_mults=(1, 2, 4, 8), # Multipliers for dimension increase at each downsampling level
                 time_emb_dim=128,  # Dimension of the time embedding
                 groups=8):         # Number of groups for Group Normalization
        super().__init__()

        self.time_mlp = nn.Sequential(
            SinusoidalPositionalEmbedding(dim),
            nn.Linear(dim, time_emb_dim),
            nn.SiLU(),
            nn.Linear(time_emb_dim, time_emb_dim)
        )

        # Initial convolution to project image channels to base dimension
        self.init_conv = nn.Conv2d(image_channels, dim, 7, padding=3)

        dims = [dim, *map(lambda m: dim * m, dim_mults)] # [dim, dim*1, dim*2, dim*4, ...]
        in_out = list(zip(dims[:-1], dims[1:])) # [(dim, dim*1), (dim*1, dim*2), ...]

        # Encoder (downsampling path)
        self.downs = nn.ModuleList([])
        for ind, (dim_in, dim_out) in enumerate(in_out):
            self.downs.append(nn.ModuleList([
                ResnetBlock(dim_in, dim_out, time_emb_dim, groups=groups),
                Downsample(dim_out) if ind != (len(in_out) - 1) else nn.Identity() # No downsample at bottleneck
            ]))

        # Bottleneck (middle layer)
        mid_dim = dims[-1]
        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)
        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)

        # Decoder (upsampling path)
        self.ups = nn.ModuleList([])
        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
            self.ups.append(nn.ModuleList([
                ResnetBlock(dim_out * 2, dim_in, time_emb_dim, groups=groups), # *2 for skip connection concatenation
                Upsample(dim_in) if ind != 0 else nn.Identity() # No upsample at final layer
            ]))

        # Final output convolution
        self.final_conv = nn.Sequential(
            Block(dim, dim, groups=groups),
            nn.Conv2d(dim, image_channels, 1) # Output noise with original image channels
        )

    def forward(self, x, time):
        # Time embedding
        t = self.time_mlp(time)

        # Initial convolution
        x = self.init_conv(x)

        h = [] # To store skip connections outputs

        # Downsampling path
        for resnet_block, downsample in self.downs:
            x = resnet_block(x, t)
            h.append(x) # Store for skip connection
            x = downsample(x)

        # Bottleneck
        x = self.mid_block1(x, t)
        x = self.mid_block2(x, t)

        # Upsampling path
        for resnet_block, upsample in self.ups:
            # Concatenate with skip connection from encoder
            x = torch.cat((x, h.pop()), dim=1)
            x = resnet_block(x, t)
            x = upsample(x)

        # Final output
        return self.final_conv(x)

Step 5: The Training Process

We will define the loss function and the training loop. The p_losses function encapsulates the core training step for a single batch.

def p_losses(denoise_model, x_start, t, noise=None):
    """
    Calculates the loss for a given batch of original images and time steps.
    
    Args:
        denoise_model (nn.Module): The U-Net model to train.
        x_start (torch.Tensor): The original, clean images (x_0).
        t (torch.Tensor): The time steps for each image in the batch.
        noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
        
    Returns:
        torch.Tensor: The mean squared error loss.
    """
    if noise is None:
        noise = torch.randn_like(x_start) # Sample noise for the forward process

    # Apply noise to get x_t
    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
    
    # Predict the noise using the denoising model
    predicted_noise = denoise_model(x_noisy, t)

    # Calculate MSE loss between actual noise and predicted noise
    loss = F.mse_loss(noise, predicted_noise)
    return loss

Step 6: Sampling (Generating New Images)

This is the reverse process, where we start from pure noise and iteratively denoise it to generate new images.

@torch.no_grad() # Disable gradient calculations for sampling
def p_sample(model, x, t, t_index):
    """
    Performs one step of the reverse diffusion process (denoising).
    Estimates x_{t-1} from x_t.
    
    Args:
        model (nn.Module): The trained U-Net model.
        x (torch.Tensor): The noisy image at time t (x_t).
        t (torch.Tensor): The current time step.
        t_index (int): The integer index of the current time step.
        
    Returns:
        torch.Tensor: The denoised image at time t-1 (x_{t-1}).
    """
    betas_t = extract(betas, t, x.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(
        sqrt_one_minus_alphas_cumprod, t, x.shape
    )
    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
    
    # Predict the noise using the model
    model_mean = sqrt_recip_alphas_t * (
        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
    )
    
    # If t is the first step (t=0), there's no more noise to add
    if t_index == 0:
        return model_mean
    else:
        posterior_variance_t = extract(posterior_variance, t, x.shape)
        # Sample new noise for the reverse step (unless t=0)
        noise = torch.randn_like(x)
        return model_mean + torch.sqrt(posterior_variance_t) * noise

@torch.no_grad()
def p_sample_loop(model, shape):
    """
    Generates a batch of images by iteratively denoising from pure noise.
    
    Args:
        model (nn.Module): The trained U-Net model.
        shape (tuple): The shape of the images to generate (batch_size, channels, height, width).
        
    Returns:
        torch.Tensor: A batch of generated images.
    """
    batch_size = shape[0]
    # Start with pure noise
    img = torch.randn(shape, device=DEVICE)

    # Iterate backwards through time steps
    for i in tqdm(reversed(range(0, TIMESTEPS)), desc='sampling loop time step', total=TIMESTEPS):
        t = torch.full((batch_size,), i, device=DEVICE, dtype=torch.long)
        img = p_sample(model, img, t, i) # Denoise one step

    return img

@torch.no_grad()
def sample_and_save_images(model, epoch, num_samples=16):
    """
    Generates and saves a grid of sample images.
    """
    # Define the shape of the images to generate
    sample_shape = (num_samples, 1, IMG_SIZE, IMG_SIZE)
    
    # Generate images
    samples = p_sample_loop(model, sample_shape)
    
    # Normalize images to [0, 1] range and convert to PIL Image
    samples = (samples + 1) * 0.5 # [-1, 1] to [0, 1]
    samples = samples.clamp(0, 1)
    
    # Create a grid of images
    grid_size = int(math.sqrt(num_samples))
    
    # Create a blank image to paste samples onto
    combined_image_width = grid_size * IMG_SIZE
    combined_image_height = grid_size * IMG_SIZE
    combined_image = Image.new('L', (combined_image_width, combined_image_height)) # 'L' for grayscale
    
    for i in range(num_samples):
        row = i // grid_size
        col = i % grid_size
        
        # Convert tensor to numpy array, scale to 0-255, convert to uint8
        img_array = (samples[i].squeeze().cpu().numpy() * 255).astype(np.uint8)
        img = Image.fromarray(img_array)
        
        # Paste onto the combined image
        combined_image.paste(img, (col * IMG_SIZE, row * IMG_SIZE))
        
    # Save the combined image
    filepath = os.path.join(SAVE_DIR, f"epoch_{epoch:04d}_samples.png")
    combined_image.save(filepath)
    print(f"Saved {num_samples} samples to {filepath}")

Putting It All Together: Training and Generation

Now we have all the pieces. We will set up the data loading, initialize the model and optimizer, and run the training loop. After training, we can use the sample_and_save_images function to generate new images.

The overall workflow during training is:

  1. Load a batch of real images (x_0).
  2. Randomly choose a time step (t) for each image in the batch.
  3. Add noise to (x_0) to get (x_t) and simultaneously record the noise (\epsilon) that was added.
  4. Feed (x_t) and (t) into the U-Net to predict the noise, (\epsilon_\theta).
  5. Calculate the MSE loss between (\epsilon) and (\epsilon_\theta).
  6. Perform backpropagation and update the U-Net's weights.

After training, to generate new images:

  1. Start with a tensor of pure random noise.
  2. Iteratively apply the p_sample function, moving backward from (T) down to 0. Each step removes a bit of noise based on the U-Net's prediction.
  3. The final output is a newly generated image.

Conclusion

Diffusion models represent a significant leap forward in generative AI. By understanding the simple yet powerful concept of gradually adding and then learning to reverse noise, developers can unlock incredible capabilities for generating highly realistic and diverse data. We have explored the forward noising process, the crucial reverse denoising process powered by U-Nets, and the key mathematical and architectural components involved.

With the conceptual understanding and the step-by-step implementation guide provided, you are now equipped to delve deeper, experiment with different schedules, model architectures, and apply these fascinating models to your own creative and technical challenges. The field is rapidly evolving, and your journey into diffusion models has just begun!

Addendum: Full Running Example Code

This section provides a complete, runnable Python script that implements the diffusion model for generating MNIST-like grayscale images. This code integrates all the snippets and additional necessary components like data loading and the main training loop.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import os
from PIL import Image
import math
from tqdm import tqdm # For progress bars

# --- Configuration Parameters ---
IMG_SIZE = 28          # Size of the input images (e.g., 28 for MNIST)
BATCH_SIZE = 128       # Number of images per training batch
NUM_EPOCHS = 100       # Number of training epochs
LEARNING_RATE = 1e-4   # Learning rate for the optimizer
TIMESTEPS = 1000       # Total number of diffusion steps (T)
BETA_START = 1e-4      # Start value for the linear beta schedule
BETA_END = 0.02        # End value for the linear beta schedule
SAVE_DIR = "diffusion_results" # Directory to save generated images
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Use GPU if available

# Ensure save directory exists
os.makedirs(SAVE_DIR, exist_ok=True)

print(f"Using device: {DEVICE}")

# --- Data Loading and Preprocessing ---
def load_mnist_data(batch_size, img_size):
    """
    Loads and preprocesses the MNIST dataset.
    Images are normalized to the range [-1, 1] for better model stability.
    """
    transform = transforms.Compose([
        transforms.Resize(img_size),
        transforms.ToTensor(), # Converts to [0, 1] range
        transforms.Normalize((0.5,), (0.5,)) # Normalizes to [-1, 1] range
    ])
    
    dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    return dataloader

# --- Variance Schedule Pre-computation ---
def linear_beta_schedule(timesteps, beta_start, beta_end):
    """
    Generates a linear schedule for beta values.
    
    Args:
        timesteps (int): The total number of diffusion steps (T).
        beta_start (float): The starting value for beta.
        beta_end (float): The ending value for beta.
        
    Returns:
        torch.Tensor: A tensor of beta values for each timestep.
    """
    return torch.linspace(beta_start, beta_end, timesteps)

# Pre-compute the schedule values and move to device
betas = linear_beta_schedule(TIMESTEPS, BETA_START, BETA_END).to(DEVICE)
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # alpha_bar_t = product(alpha_s from s=1 to t)
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) # alpha_bar_{t-1}
sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # Used in reverse process
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) # sqrt(alpha_bar_t)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) # sqrt(1 - alpha_bar_t)
# Variance for the reverse step, used in p_sample
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)

# --- Forward Diffusion Helper Functions ---
def extract(a, t, x_shape):
    """
    Extracts the values from a tensor 'a' at given indices 't' and reshapes them
    to match the shape of 'x_shape'. This is used to select the correct
    alpha_bar_t or sqrt(1 - alpha_bar_t) for each sample in a batch.
    
    Args:
        a (torch.Tensor): The tensor to extract values from (e.g., alphas_cumprod).
        t (torch.Tensor): A tensor of time indices for each sample in the batch.
        x_shape (torch.Size): The desired shape for the extracted values.
        
    Returns:
        torch.Tensor: The extracted and reshaped tensor.
    """
    batch_size = t.shape[0]
    out = a.gather(-1, t.cpu()) # Extract values using cpu indices
    # Reshape to (batch_size, 1, 1, 1) for broadcasting with image tensors
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def q_sample(x_start, t, noise=None):
    """
    Applies noise to the original image x_start at time step t.
    This implements the forward diffusion process using the reparameterization trick.
    
    Args:
        x_start (torch.Tensor): The original, clean image (x_0).
        t (torch.Tensor): The current time step for each sample in the batch.
        noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
        
    Returns:
        torch.Tensor: The noisy image x_t.
    """
    if noise is None:
        noise = torch.randn_like(x_start) # Sample Gaussian noise

    # Extract sqrt(alpha_bar_t) and sqrt(1 - alpha_bar_t) for the current batch and time steps
    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape)

    # Apply the reparameterization trick: x_t = sqrt(alpha_bar_t)*x_0 + sqrt(1 - alpha_bar_t)*epsilon
    x_t = sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
    return x_t

# --- Denoising U-Net Model Definition ---
class SinusoidalPositionalEmbedding(nn.Module):
    """
    Generates sinusoidal positional embeddings for time steps.
    These embeddings allow the model to understand the current time step (noise level).
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        # Compute the sinusoidal arguments
        embeddings = math.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings

class Block(nn.Module):
    """
    A basic convolutional block with Group Normalization and SiLU (Swish) activation.
    Includes an optional scale and shift for time embedding integration.
    """
    def __init__(self, dim_in, dim_out, groups=8):
        super().__init__()
        self.proj = nn.Conv2d(dim_in, dim_out, 3, padding=1)
        self.norm = nn.GroupNorm(groups, dim_out)
        self.act = nn.SiLU()

    def forward(self, x, scale_shift=None):
        x = self.proj(x)
        x = self.norm(x)
        
        # Apply optional scale and shift from time embedding
        if scale_shift is not None:
            scale, shift = scale_shift
            x = x * (scale + 1) + shift

        x = self.act(x)
        return x

class ResnetBlock(nn.Module):
    """
    A Residual Block, commonly used in U-Nets.
    It includes two convolutional blocks and a skip connection.
    Time embeddings are incorporated here by modulating the normalization layer.
    """
    def __init__(self, dim_in, dim_out, time_emb_dim=None, groups=8):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim, dim_out * 2)
        ) if time_emb_dim is not None else None

        self.block1 = Block(dim_in, dim_out, groups=groups)
        self.block2 = Block(dim_out, dim_out, groups=groups)
        self.res_conv = nn.Conv2d(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity()

    def forward(self, x, time_emb=None):
        scale_shift = None
        if self.mlp is not None and time_emb is not None:
            time_emb = self.mlp(time_emb)
            # Split the output into scale and shift for normalization
            time_emb = time_emb.reshape(time_emb.shape[0], time_emb.shape[1], 1, 1)
            scale_shift = time_emb.chunk(2, dim=1) # Split into two along dimension 1

        h = self.block1(x, scale_shift=scale_shift)
        h = self.block2(h)
        return h + self.res_conv(x) # Add residual connection

class Downsample(nn.Module):
    """
    Downsampling layer using a stride-2 convolution.
    """
    def __init__(self, dim):
        super().__init__()
        self.conv = nn.Conv2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1

    def forward(self, x):
        return self.conv(x)

class Upsample(nn.Module):
    """
    Upsampling layer using a transposed convolution.
    """
    def __init__(self, dim):
        super().__init__()
        self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1) # 4x4 kernel, stride 2, padding 1

    def forward(self, x):
        return self.conv(x)

class Unet(nn.Module):
    """
    The Denoising U-Net model.
    It takes a noisy image and a time step, and outputs the predicted noise.
    """
    def __init__(self,
                 dim,               # Base dimension for the model
                 image_channels=1,  # Number of channels in the input image (e.g., 1 for grayscale)
                 dim_mults=(1, 2, 4, 8), # Multipliers for dimension increase at each downsampling level
                 time_emb_dim=128,  # Dimension of the time embedding
                 groups=8):         # Number of groups for Group Normalization
        super().__init__()

        # Time embedding MLP
        self.time_mlp = nn.Sequential(
            SinusoidalPositionalEmbedding(dim),
            nn.Linear(dim, time_emb_dim),
            nn.SiLU(),
            nn.Linear(time_emb_dim, time_emb_dim)
        )

        # Initial convolution to project image channels to base dimension
        self.init_conv = nn.Conv2d(image_channels, dim, 7, padding=3)

        # Calculate dimensions for encoder/decoder blocks
        dims = [dim, *map(lambda m: dim * m, dim_mults)] # [dim, dim*1, dim*2, dim*4, ...]
        in_out = list(zip(dims[:-1], dims[1:])) # [(dim, dim*1), (dim*1, dim*2), ...]

        # Encoder (downsampling path)
        self.downs = nn.ModuleList([])
        for ind, (dim_in, dim_out) in enumerate(in_out):
            self.downs.append(nn.ModuleList([
                ResnetBlock(dim_in, dim_out, time_emb_dim, groups=groups),
                Downsample(dim_out) if ind != (len(in_out) - 1) else nn.Identity() # No downsample at bottleneck
            ]))

        # Bottleneck (middle layer)
        mid_dim = dims[-1]
        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)
        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim, groups=groups)

        # Decoder (upsampling path)
        self.ups = nn.ModuleList([])
        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
            self.ups.append(nn.ModuleList([
                ResnetBlock(dim_out * 2, dim_in, time_emb_dim, groups=groups), # *2 for skip connection concatenation
                Upsample(dim_in) if ind != 0 else nn.Identity() # No upsample at final layer
            ]))

        # Final output convolution
        self.final_conv = nn.Sequential(
            Block(dim, dim, groups=groups),
            nn.Conv2d(dim, image_channels, 1) # Output noise with original image channels
        )

    def forward(self, x, time):
        # Time embedding
        t = self.time_mlp(time)

        # Initial convolution
        x = self.init_conv(x)

        h = [] # To store skip connections outputs for concatenation

        # Downsampling path
        for resnet_block, downsample in self.downs:
            x = resnet_block(x, t)
            h.append(x) # Store for skip connection
            x = downsample(x)

        # Bottleneck
        x = self.mid_block1(x, t)
        x = self.mid_block2(x, t)

        # Upsampling path
        for resnet_block, upsample in self.ups:
            # Concatenate with skip connection from encoder
            x = torch.cat((x, h.pop()), dim=1)
            x = resnet_block(x, t)
            x = upsample(x)

        # Final output
        return self.final_conv(x)

# --- Loss Function for Training ---
def p_losses(denoise_model, x_start, t, noise=None):
    """
    Calculates the loss for a given batch of original images and time steps.
    
    Args:
        denoise_model (nn.Module): The U-Net model to train.
        x_start (torch.Tensor): The original, clean images (x_0).
        t (torch.Tensor): The time steps for each image in the batch.
        noise (torch.Tensor, optional): Pre-sampled noise. If None, noise is sampled.
        
    Returns:
        torch.Tensor: The mean squared error loss.
    """
    if noise is None:
        noise = torch.randn_like(x_start) # Sample noise for the forward process

    # Apply noise to get x_t
    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
    
    # Predict the noise using the denoising model
    predicted_noise = denoise_model(x_noisy, t)

    # Calculate MSE loss between actual noise and predicted noise
    loss = F.mse_loss(noise, predicted_noise)
    return loss

# --- Sampling (Image Generation) Functions ---
@torch.no_grad() # Disable gradient calculations for sampling
def p_sample(model, x, t, t_index):
    """
    Performs one step of the reverse diffusion process (denoising).
    Estimates x_{t-1} from x_t.
    
    Args:
        model (nn.Module): The trained U-Net model.
        x (torch.Tensor): The noisy image at time t (x_t).
        t (torch.Tensor): The current time step.
        t_index (int): The integer index of the current time step.
        
    Returns:
        torch.Tensor: The denoised image at time t-1 (x_{t-1}).
    """
    betas_t = extract(betas, t, x.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(
        sqrt_one_minus_alphas_cumprod, t, x.shape
    )
    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
    
    # Predict the noise using the model
    # This formula is derived from the reverse process mean estimation
    model_mean = sqrt_recip_alphas_t * (
        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
    )
    
    # If t is the first step (t=0), there's no more noise to add
    if t_index == 0:
        return model_mean
    else:
        # Add noise sampled from the posterior distribution's variance
        posterior_variance_t = extract(posterior_variance, t, x.shape)
        noise = torch.randn_like(x)
        return model_mean + torch.sqrt(posterior_variance_t) * noise

@torch.no_grad()
def p_sample_loop(model, shape):
    """
    Generates a batch of images by iteratively denoising from pure noise.
    
    Args:
        model (nn.Module): The trained U-Net model.
        shape (tuple): The shape of the images to generate (batch_size, channels, height, width).
        
    Returns:
        torch.Tensor: A batch of generated images.
    """
    batch_size = shape[0]
    # Start with pure Gaussian noise
    img = torch.randn(shape, device=DEVICE)

    # Iterate backwards through time steps, denoising at each step
    for i in tqdm(reversed(range(0, TIMESTEPS)), desc='sampling loop time step', total=TIMESTEPS):
        t = torch.full((batch_size,), i, device=DEVICE, dtype=torch.long)
        img = p_sample(model, img, t, i) # Denoise one step

    return img

@torch.no_grad()
def sample_and_save_images(model, epoch, num_samples=16):
    """
    Generates and saves a grid of sample images.
    
    Args:
        model (nn.Module): The trained U-Net model.
        epoch (int): The current epoch number, used for naming the saved file.
        num_samples (int): The number of images to generate and save.
    """
    # Define the shape of the images to generate (e.g., 16 samples, 1 channel, 28x28 pixels)
    sample_shape = (num_samples, 1, IMG_SIZE, IMG_SIZE)
    
    # Generate images using the reverse diffusion process
    samples = p_sample_loop(model, sample_shape)
    
    # Normalize images from [-1, 1] to [0, 1] range for saving as image files
    samples = (samples + 1) * 0.5 
    samples = samples.clamp(0, 1) # Ensure values are within valid range
    
    # Create a grid of images for visualization
    grid_size = int(math.sqrt(num_samples))
    
    # Create a blank image to paste samples onto
    combined_image_width = grid_size * IMG_SIZE
    combined_image_height = grid_size * IMG_SIZE
    # 'L' mode for grayscale images
    combined_image = Image.new('L', (combined_image_width, combined_image_height)) 
    
    for i in range(num_samples):
        row = i // grid_size
        col = i % grid_size
        
        # Convert tensor to numpy array, scale to 0-255, convert to uint8
        img_array = (samples[i].squeeze().cpu().numpy() * 255).astype(np.uint8)
        img = Image.fromarray(img_array)
        
        # Paste the individual generated image onto the combined grid image
        combined_image.paste(img, (col * IMG_SIZE, row * IMG_SIZE))
        
    # Save the combined image to the specified directory
    filepath = os.path.join(SAVE_DIR, f"epoch_{epoch:04d}_samples.png")
    combined_image.save(filepath)
    print(f"Saved {num_samples} samples to {filepath}")


# --- Main Training Loop ---
def train_diffusion_model():
    """
    Main function to train the diffusion model.
    Initializes the model, optimizer, loads data, and runs the training epochs.
    """
    # Load data
    dataloader = load_mnist_data(BATCH_SIZE, IMG_SIZE)

    # Initialize model
    model = Unet(
        dim=64, # Base dimension for the U-Net
        image_channels=1, # MNIST is grayscale
        dim_mults=(1, 2, 4), # Dimension multipliers for downsampling path
        time_emb_dim=256 # Dimension of time embeddings
    ).to(DEVICE)
    
    # Initialize optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

    print("Starting training...")
    for epoch in range(NUM_EPOCHS):
        model.train() # Set model to training mode
        total_loss = 0
        for step, (images, _) in enumerate(tqdm(dataloader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS}")):
            optimizer.zero_grad() # Clear gradients

            images = images.to(DEVICE)
            
            # Sample random time steps for the current batch
            t = torch.randint(0, TIMESTEPS, (images.shape[0],), device=DEVICE).long()

            # Calculate loss
            loss = p_losses(model, images, t)
            total_loss += loss.item()

            # Backpropagation and optimization step
            loss.backward()
            optimizer.step()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1} completed. Average Loss: {avg_loss:.4f}")

        # Generate and save sample images periodically
        if (epoch + 1) % 10 == 0 or epoch == 0: # Save samples at epoch 0 and every 10 epochs
            model.eval() # Set model to evaluation mode for sampling
            sample_and_save_images(model, epoch + 1, num_samples=16)

    print("Training finished.")
    # Optionally save the final model
    torch.save(model.state_dict(), os.path.join(SAVE_DIR, "diffusion_model_final.pth"))
    print(f"Final model saved to {os.path.join(SAVE_DIR, 'diffusion_model_final.pth')}")

if __name__ == "__main__":
    train_diffusion_model()