breakpilot-compliance/ai-compliance-sdk/internal/iace/tech_file_generator.go

package iace

import (
	"context"
	"fmt"
	"strings"

	"github.com/breakpilot/ai-compliance-sdk/internal/llm"
	"github.com/breakpilot/ai-compliance-sdk/internal/ucca"
	"github.com/google/uuid"
)

// ============================================================================
// TechFileGenerator — LLM-based generation of technical file sections
// ============================================================================

// TechFileGenerator generates technical file section content using LLM and RAG.
type TechFileGenerator struct {
	llmRegistry *llm.ProviderRegistry
	ragClient   *ucca.LegalRAGClient
	store       *Store
}

// NewTechFileGenerator creates a new TechFileGenerator.
func NewTechFileGenerator(registry *llm.ProviderRegistry, ragClient *ucca.LegalRAGClient, store *Store) *TechFileGenerator {
	return &TechFileGenerator{
		llmRegistry: registry,
		ragClient:   ragClient,
		store:       store,
	}
}

// SectionGenerationContext holds all project data needed for LLM section generation.
type SectionGenerationContext struct {
	Project         *Project
	Components      []Component
	Hazards         []Hazard
	Assessments     map[uuid.UUID][]RiskAssessment // keyed by hazardID
	Mitigations     map[uuid.UUID][]Mitigation     // keyed by hazardID
	Classifications []RegulatoryClassification
	Evidence        []Evidence
	RAGContext      string // aggregated text from RAG search
}

// ============================================================================
// BuildSectionContext — loads all project data + RAG context
// ============================================================================

// BuildSectionContext loads project data and RAG context for a given section type.
func (g *TechFileGenerator) BuildSectionContext(ctx context.Context, projectID uuid.UUID, sectionType string) (*SectionGenerationContext, error) {
	// Load project
	project, err := g.store.GetProject(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load project: %w", err)
	}
	if project == nil {
		return nil, fmt.Errorf("project %s not found", projectID)
	}

	// Load components
	components, err := g.store.ListComponents(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load components: %w", err)
	}

	// Load hazards
	hazards, err := g.store.ListHazards(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load hazards: %w", err)
	}

	// Load assessments and mitigations per hazard
	assessments := make(map[uuid.UUID][]RiskAssessment)
	mitigations := make(map[uuid.UUID][]Mitigation)
	for _, h := range hazards {
		a, err := g.store.ListAssessments(ctx, h.ID)
		if err != nil {
			return nil, fmt.Errorf("load assessments for hazard %s: %w", h.ID, err)
		}
		assessments[h.ID] = a

		m, err := g.store.ListMitigations(ctx, h.ID)
		if err != nil {
			return nil, fmt.Errorf("load mitigations for hazard %s: %w", h.ID, err)
		}
		mitigations[h.ID] = m
	}

	// Load classifications
	classifications, err := g.store.GetClassifications(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load classifications: %w", err)
	}

	// Load evidence
	evidence, err := g.store.ListEvidence(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("load evidence: %w", err)
	}

	// Perform RAG search for section-specific context
	ragContext := ""
	if g.ragClient != nil {
		ragQuery := buildRAGQuery(sectionType)
		results, ragErr := g.ragClient.SearchCollection(ctx, "bp_iace_libraries", ragQuery, nil, 5)
		if ragErr == nil && len(results) > 0 {
			var ragParts []string
			for _, r := range results {
				entry := fmt.Sprintf("[%s] %s", r.RegulationShort, truncateForPrompt(r.Text, 400))
				ragParts = append(ragParts, entry)
			}
			ragContext = strings.Join(ragParts, "\n\n")
		}
		// RAG failure is non-fatal — we proceed without context
	}

	return &SectionGenerationContext{
		Project:         project,
		Components:      components,
		Hazards:         hazards,
		Assessments:     assessments,
		Mitigations:     mitigations,
		Classifications: classifications,
		Evidence:        evidence,
		RAGContext:      ragContext,
	}, nil
}

// ============================================================================
// GenerateSection — main entry point
// ============================================================================

// GenerateSection generates the content for a technical file section using LLM.
// If LLM is unavailable, returns an enhanced placeholder with project data.
func (g *TechFileGenerator) GenerateSection(ctx context.Context, projectID uuid.UUID, sectionType string) (string, error) {
	sctx, err := g.BuildSectionContext(ctx, projectID, sectionType)
	if err != nil {
		return "", fmt.Errorf("build section context: %w", err)
	}

	// Build prompts
	systemPrompt := getSystemPrompt(sectionType)
	userPrompt := buildUserPrompt(sctx, sectionType)

	// Attempt LLM generation
	resp, err := g.llmRegistry.Chat(ctx, &llm.ChatRequest{
		Messages: []llm.Message{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userPrompt},
		},
		Temperature: 0.15,
		MaxTokens:   4096,
	})
	if err != nil {
		// LLM unavailable — return structured fallback with real project data
		return buildFallbackContent(sctx, sectionType), nil
	}

	return resp.Message.Content, nil
}