breakpilot-compliance/ai-compliance-sdk/internal/ucca/legal_rag_client.go

package ucca

import (
	"bytes"
	"context"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"
)

// LegalRAGClient provides access to the compliance CE vector search via Qdrant + Ollama bge-m3.
type LegalRAGClient struct {
	qdrantURL        string
	qdrantAPIKey     string
	ollamaURL        string
	embeddingModel   string
	collection       string
	httpClient       *http.Client
	textIndexEnsured map[string]bool
	hybridEnabled    bool
	graphEnabled     bool

	// Blue-Green „authoritative slice promotion" (additiv, KEIN CE-Ersatz): faellt eine Query
	// in den KB-2026.1-Scope (DP/CRA/MaschVO/NIS2/DataAct/DORA/AIAct + EDPB/DSK-Guidance), wird
	// die hochwertige Slice-Collection abgefragt; sonst bleibt der breite Default (bp_compliance_ce).
	kbSliceCollection     string
	kbScopeRoutingEnabled bool
}

// NewLegalRAGClient creates a new Legal RAG client using Ollama bge-m3 embeddings.
func NewLegalRAGClient() *LegalRAGClient {
	qdrantURL := os.Getenv("QDRANT_URL")
	if qdrantURL == "" {
		qdrantURL = "http://localhost:6333"
	}
	qdrantURL = strings.TrimRight(qdrantURL, "/")

	qdrantAPIKey := os.Getenv("QDRANT_API_KEY")

	ollamaURL := os.Getenv("OLLAMA_URL")
	if ollamaURL == "" {
		ollamaURL = "http://localhost:11434"
	}

	hybridEnabled := os.Getenv("RAG_HYBRID_SEARCH") != "false"
	// Graph-Expansion ist OPT-IN: kein gemessener Rang-Nutzen ggue. der Binding-Augmentation,
	// +1 Qdrant-Call/Suche, Flutungsrisiko ueber Reverse-Kanten. Bleibt als Recall-Sicherheitsnetz
	// fuer spaetere Luecken (RAG_GRAPH_EXPANSION=true). Die Graph-Kanten werden in der Response
	// zur Begruendung/Vollstaendigkeit genutzt, nicht zur Pool-Expansion (Default).
	graphEnabled := os.Getenv("RAG_GRAPH_EXPANSION") == "true"

	// KB-2026.1 authoritative slice (Blue-Green, additiv). Routing default AN; Rollback ohne
	// Redeploy ueber RAG_KB_SCOPE_ROUTING=false (dann faellt alles auf den CE-Default zurueck).
	kbSlice := os.Getenv("RAG_KB_SLICE_COLLECTION")
	if kbSlice == "" {
		kbSlice = "kb_2026_1_build"
	}
	kbScopeRouting := os.Getenv("RAG_KB_SCOPE_ROUTING") != "false"

	return &LegalRAGClient{
		qdrantURL:             qdrantURL,
		qdrantAPIKey:          qdrantAPIKey,
		ollamaURL:             ollamaURL,
		embeddingModel:        "bge-m3",
		collection:            "bp_compliance_ce",
		textIndexEnsured:      make(map[string]bool),
		hybridEnabled:         hybridEnabled,
		graphEnabled:          graphEnabled,
		kbSliceCollection:     kbSlice,
		kbScopeRoutingEnabled: kbScopeRouting,
		httpClient: &http.Client{
			Timeout: 60 * time.Second,
		},
	}
}

// SearchCollection queries a specific Qdrant collection for relevant passages.
// If collection is empty, it falls back to the default collection (bp_compliance_ce).
func (c *LegalRAGClient) SearchCollection(ctx context.Context, collection string, query string, regulationIDs []string, topK int) ([]LegalSearchResult, error) {
	return c.searchInternal(ctx, c.resolveCollection(query, collection), query, regulationIDs, topK)
}

// Search queries the compliance corpus for relevant passages. The target collection is resolved by
// the Blue-Green slice routing: the KB-2026.1 slice for in-scope queries, else the broad CE default.
func (c *LegalRAGClient) Search(ctx context.Context, query string, regulationIDs []string, topK int) ([]LegalSearchResult, error) {
	return c.searchInternal(ctx, c.resolveCollection(query, ""), query, regulationIDs, topK)
}

// searchInternal performs the actual search against a given collection.
// If hybrid search is enabled, it uses the Qdrant Query API with RRF fusion
// (dense + full-text). Falls back to dense-only /points/search on failure.
func (c *LegalRAGClient) searchInternal(ctx context.Context, collection string, query string, regulationIDs []string, topK int) ([]LegalSearchResult, error) {
	// Multi-Regulation-Retrieval: nennt die Query EXPLIZIT >=2 Regelwerke (z.B. "CRA und
	// Maschinenverordnung"), wird pro Regelwerk separat retrieved + gemergt, damit BEIDE
	// Domaenen im Prompt landen statt nur der keyword-dominanten. Generisch (Query->Regelwerke,
	// keine doc-spezifische Logik); nur wenn der Caller nicht ohnehin schon auf Regulierungen
	// filtert. Best-effort: leeres/fehlerhaftes Multi-Ergebnis faellt auf die Standardsuche zurueck.
	if len(regulationIDs) == 0 {
		if regs := detectRegulations(query); len(regs) >= 2 {
			if mr, mErr := c.searchMultiRegulation(ctx, collection, query, regs, topK); mErr == nil && len(mr) > 0 {
				return mr, nil
			}
		}
	}

	embedding, err := c.generateEmbedding(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("failed to generate embedding: %w", err)
	}

	var hits []qdrantSearchHit

	if c.hybridEnabled {
		hybridHits, err := c.searchHybrid(ctx, collection, embedding, regulationIDs, topK)
		if err == nil {
			hits = hybridHits
		}
	}

	if hits == nil {
		denseHits, err := c.searchDense(ctx, collection, embedding, regulationIDs, topK)
		if err != nil {
			return nil, err
		}
		hits = denseHits
	}

	// Stratified: den binding_law-Pool ERGAENZEN (nicht ersetzen), damit die Pflichtquelle
	// immer Kandidat ist — Guidance bleibt als Auslegungskontext erhalten. Best-effort:
	// Fehler beim Binding-Query degradieren still auf den semantischen Pool.
	if bindingHits, bErr := c.searchBinding(ctx, collection, embedding, topK); bErr == nil {
		hits = mergeDedupHits(hits, bindingHits)
	}

	// Control-Augmentation: bei expliziter Umsetzungsfrage einen tiefen dense-Pool ziehen und
	// nur die Control-Pool-Rollen behalten — so werden NIST/CRA-Anhang (dense rank ~8-9, unter
	// dem kleinen top-K) Kandidaten. Re-Rank/applyControlRoles ordnen sie danach.
	if queryWantsControls(query) {
		if controlHits, cErr := c.searchControls(ctx, collection, embedding); cErr == nil {
			hits = mergeDedupHits(hits, controlHits)
		}
	}

	// Graph-Augmentation: verbundene Normen (references_out/in) der Top-Hits ueber die
	// praezise Zitations-Kante in den Pool ziehen — z.B. Art. 13 CRA zieht Anhang I (die
	// eigentliche Pflichtquelle). Pool-Augmentation only; Re-Rank + topK bleiben.
	if c.graphEnabled {
		hits = c.expandViaGraph(ctx, collection, hits)
	}

	results := hitsToResults(hits)

	// Authority-aware Re-Ranking: bindendes Recht der passenden Jurisdiktion/Domaene nach
	// oben, Guidance/Fremdrecht/Off-Domain runter (nichts wird geloescht). Reihenfolge only,
	// Response-Schema unveraendert. Score traegt den Authority-Score, damit nachgelagerte
	// Multi-Collection-Merges (Advisor) die Ordnung bewahren.
	results = rerankByAuthority(query, results)

	// Control-Diversity: auf einer Umsetzungsfrage darf impl_guidance (ENISA) Top-1 bleiben,
	// aber die Top-K soll mindestens eine binding operational_requirement (CRA Anhang I) und
	// einen control_standard (NIST/ISO) zeigen, falls im Pool — Quellenarten sichtbar machen
	// statt sie kuenstlich auf Top-1 zu heben. Nur Reihenfolge, vor der Truncation.
	if queryWantsControls(query) {
		results = ensureControlDiversity(results, topK)
	}

	if topK > 0 && len(results) > topK {
		results = results[:topK]
	}

	return results, nil
}

// mergeDedupHits concatenates two hit lists, keeping the first occurrence of each point ID.
func mergeDedupHits(primary, extra []qdrantSearchHit) []qdrantSearchHit {
	seen := make(map[string]bool, len(primary)+len(extra))
	out := make([]qdrantSearchHit, 0, len(primary)+len(extra))
	for _, list := range [][]qdrantSearchHit{primary, extra} {
		for _, h := range list {
			id := fmt.Sprint(h.ID)
			if seen[id] {
				continue
			}
			seen[id] = true
			out = append(out, h)
		}
	}
	return out
}

// FormatLegalContextForPrompt formats the legal context for inclusion in an LLM prompt.
func (c *LegalRAGClient) FormatLegalContextForPrompt(lc *LegalContext) string {
	if lc == nil || len(lc.Results) == 0 {
		return ""
	}

	var buf bytes.Buffer
	buf.WriteString("\n\n**Relevante Rechtsgrundlagen:**\n\n")

	for i, result := range lc.Results {
		buf.WriteString(fmt.Sprintf("%d. **%s** (%s)", i+1, result.RegulationShort, result.RegulationCode))
		if len(result.Pages) > 0 {
			buf.WriteString(fmt.Sprintf(" - Seiten %v", result.Pages))
		}
		buf.WriteString("\n")
		buf.WriteString(fmt.Sprintf("   > %s\n\n", truncateText(result.Text, 300)))
	}

	return buf.String()
}

// ListAvailableRegulations returns the list of regulations available in the corpus.
func (c *LegalRAGClient) ListAvailableRegulations() []CERegulationInfo {
	return []CERegulationInfo{
		{ID: "eu_2023_1230", NameDE: "EU-Maschinenverordnung 2023/1230", NameEN: "EU Machinery Regulation 2023/1230", Short: "Maschinenverordnung", Category: "regulation"},
		{ID: "eu_2024_1689", NameDE: "EU KI-Verordnung (AI Act)", NameEN: "EU AI Act 2024/1689", Short: "AI Act", Category: "regulation"},
		{ID: "eu_2024_2847", NameDE: "Cyber Resilience Act", NameEN: "Cyber Resilience Act 2024/2847", Short: "CRA", Category: "regulation"},
		{ID: "eu_2022_2555", NameDE: "NIS-2-Richtlinie", NameEN: "NIS2 Directive 2022/2555", Short: "NIS2", Category: "regulation"},
		{ID: "eu_2016_679", NameDE: "Datenschutz-Grundverordnung (DSGVO)", NameEN: "General Data Protection Regulation (GDPR)", Short: "DSGVO/GDPR", Category: "regulation"},
		{ID: "eu_blue_guide_2022", NameDE: "EU Blue Guide 2022", NameEN: "EU Blue Guide 2022", Short: "Blue Guide", Category: "guidance"},
		{ID: "nist_sp_800_218", NameDE: "NIST Secure Software Development Framework", NameEN: "NIST SSDF SP 800-218", Short: "NIST SSDF", Category: "guidance"},
		{ID: "nist_csf_2_0", NameDE: "NIST Cybersecurity Framework 2.0", NameEN: "NIST CSF 2.0", Short: "NIST CSF", Category: "guidance"},
		{ID: "oecd_ai_principles", NameDE: "OECD Empfehlung zu Kuenstlicher Intelligenz", NameEN: "OECD Recommendation on AI", Short: "OECD AI", Category: "guidance"},
		{ID: "enisa_supply_chain_good_practices", NameDE: "ENISA Supply Chain Cybersecurity", NameEN: "ENISA Good Practices for Supply Chain Cybersecurity", Short: "ENISA Supply Chain", Category: "guidance"},
		{ID: "enisa_threat_landscape_supply_chain", NameDE: "ENISA Threat Landscape Supply Chain", NameEN: "ENISA Threat Landscape for Supply Chain Attacks", Short: "ENISA Threat SC", Category: "guidance"},
		{ID: "enisa_ics_scada_dependencies", NameDE: "ENISA ICS/SCADA Abhaengigkeiten", NameEN: "ENISA ICS/SCADA Communication Dependencies", Short: "ENISA ICS/SCADA", Category: "guidance"},
		{ID: "cisa_secure_by_design", NameDE: "CISA Secure by Design", NameEN: "CISA Secure by Design", Short: "CISA SbD", Category: "guidance"},
		{ID: "enisa_cybersecurity_state_2024", NameDE: "ENISA State of Cybersecurity 2024", NameEN: "ENISA State of Cybersecurity in the Union 2024", Short: "ENISA 2024", Category: "guidance"},
		// BAuA — Technische Regeln (gemeinfrei, §5 UrhG)
		{ID: "trbs", NameDE: "TRBS — Technische Regeln fuer Betriebssicherheit", NameEN: "TRBS — Technical Rules for Operational Safety", Short: "TRBS", Category: "trbs"},
		{ID: "trgs", NameDE: "TRGS — Technische Regeln fuer Gefahrstoffe", NameEN: "TRGS — Technical Rules for Hazardous Substances", Short: "TRGS", Category: "trgs"},
		{ID: "asr", NameDE: "ASR — Arbeitsstaettenregeln", NameEN: "ASR — Workplace Rules", Short: "ASR", Category: "asr"},
		// OSHA
		{ID: "osha_1910", NameDE: "OSHA 1910 Subpart O — Maschinenschutz", NameEN: "OSHA 1910 Subpart O — Machinery and Machine Guarding", Short: "OSHA 1910", Category: "osha"},
		// EuGH
		{ID: "eugh_c_588_21", NameDE: "EuGH C-588/21 P — Datenschutz-Urteil", NameEN: "ECJ C-588/21 P — Data Protection Judgment", Short: "EuGH C-588/21", Category: "eu_recht"},
	}
}