agent-management/backend/src/main/java/com/agent/management/rag/graph/Neo4jGraphRagCypherGenerationService.java

package com.agent.management.rag.graph;

import com.agent.management.rag.bridge.RagAiBridgeClient;
import com.agent.management.rag.capability.RagCapabilityProfileService;
import com.agent.management.rag.capability.RagSchemaLinker;
import com.agent.management.rag.capability.RagCapabilityProfile;
import com.agent.management.rag.capability.RagSemanticCatalogService;
import com.agent.management.rag.capability.GraphBusinessSubgraphSelector;
import com.agent.management.rag.capability.RagEntityMentionExtractor;
import com.agent.management.rag.memory.RagExampleMemoryService;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.rag.model.RagSourceType;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.*;

@Service
@RequiredArgsConstructor
public class Neo4jGraphRagCypherGenerationService implements CypherGenerationService {
    private final RagAiBridgeClient bridge;
    private final RagCapabilityProfileService profiles;
    private final RagSchemaLinker linker;
    private final RagSemanticCatalogService semanticCatalog;
    private final GraphBusinessSubgraphSelector subgraphs;
    private final GraphSchemaValidator validator;
    private final RagExampleMemoryService examples;
    private final RagEntityMentionExtractor entities;

    public Optional<String> generateCypher(RagQuery query, Long graphSourceId) {
        if (!allowed(query)) return Optional.empty();
        RagCapabilityProfile profile = profiles.get(RagSourceType.GRAPH, String.valueOf(graphSourceId));
        GraphSchemaSnapshot actual = profile.graphSchema();
        GraphSchemaSnapshot authorized = actual.filter(stringList(query,"allowedLabels"),
                stringList(query,"allowedRelationships"), properties(query));
        List<Map<String,Object>> promptExamples = new ArrayList<>(maps(query, "examples"));
        promptExamples.addAll(examples.findSimilar(RagSourceType.GRAPH, String.valueOf(graphSourceId), query.getQuery(), 3));
        List<String> semanticLabels = semanticCatalog.rank(query.getQuery(), profile, "label:", 8);
        List<String> semanticRelationships = semanticCatalog.rank(query.getQuery(), profile, "relationship:", 8);
        GraphSchemaSnapshot semanticSchema = authorized.nodes().size() <= 12 || (semanticLabels.isEmpty() && semanticRelationships.isEmpty()) ? authorized
                : subgraphs.select(query.getQuery(), profile, authorized, 10, 14);
        GraphSchemaSnapshot selected = linker.selectGraphSchema(query.getQuery(), semanticSchema, promptExamples);
        if (selected.nodes().isEmpty()) throw new IllegalArgumentException("no relevant authorized graph schema is available for Text-to-Cypher");
        int maxDepth = integer(query, "maxDepth", 3);

        Map<String,Object> request = new LinkedHashMap<>();
        request.put("query",query.getQuery()); request.put("graphSourceId",graphSourceId);
        request.put("schema",selected.toPromptText()); request.put("examples",normalizeExamples(promptExamples));
        request.put("businessRules",generationContext(query)); request.put("allowedLabels",selected.labelNames());
        request.put("allowedRelationships",selected.relationshipTypeNames()); request.put("allowedProperties",allowedProperties(selected));
        request.put("maxDepth",maxDepth);
        request.put("entityMentions", entities.extract(query.getQuery()));
        Object cypher = bridge.textToCypher(request).get("cypher");
        if(cypher==null||String.valueOf(cypher).isBlank()) return Optional.empty();
        String generated=String.valueOf(cypher);
        try {
            validator.validate(generated,selected,maxDepth);
            return Optional.of(generated);
        } catch (IllegalArgumentException validationError) {
            Object fixed=bridge.repair(Map.of("language","CYPHER","question",query.getQuery(),"query",generated,
                    "error",validationError.getMessage(),"schemaText",selected.toPromptText(),"maxRows",50,"maxDepth",maxDepth)).get("query");
            if(fixed==null||String.valueOf(fixed).isBlank()) throw validationError;
            String repaired=String.valueOf(fixed);validator.validate(repaired,selected,maxDepth);return Optional.of(repaired);
        }
    }

    public Optional<String> repair(RagQuery query,Long graphSourceId,String failedCypher,String error){
        GraphSchemaSnapshot schema=profiles.get(RagSourceType.GRAPH,String.valueOf(graphSourceId)).graphSchema()
                .filter(stringList(query,"allowedLabels"),stringList(query,"allowedRelationships"),properties(query));
        int maxDepth=integer(query,"maxDepth",3);
        Object fixed=bridge.repair(Map.of("language","CYPHER","question",query.getQuery(),"query",failedCypher,
                "error",error,"schemaText",schema.toPromptText(),"maxRows",50,"maxDepth",maxDepth)).get("query");
        if(fixed==null||String.valueOf(fixed).isBlank())return Optional.empty();
        validator.validate(String.valueOf(fixed),schema,maxDepth);return Optional.of(String.valueOf(fixed));
    }

    public void validateExplicit(RagQuery query, Long graphSourceId, String cypher) {
        GraphSchemaSnapshot authorized = profiles.get(RagSourceType.GRAPH, String.valueOf(graphSourceId)).graphSchema()
                .filter(stringList(query, "allowedLabels"), stringList(query, "allowedRelationships"), properties(query));
        validator.validate(cypher, authorized, integer(query, "maxDepth", 3));
    }

    static boolean allowed(RagQuery q){if(q.getFilters()==null)return false;return Boolean.TRUE.equals(q.getFilters().get("allowTextToCypher"))||"AUTO_GENERATE".equals(String.valueOf(q.getFilters().get("retrievalMode")));}
    private static int integer(RagQuery q,String key,int d){Object v=q.getFilters().get(key);return v instanceof Number n?n.intValue():d;}
    private static String string(RagQuery q,String key){Object value=q.getFilters().get(key);return value==null?"":String.valueOf(value);}
    private static String generationContext(RagQuery query){String existing=string(query,"businessRules");List<String> rules=stringList(query,"generationRules");return rules.isEmpty()?existing:existing+(existing.isBlank()?"":"\n")+"Generation rules:\n- "+String.join("\n- ",rules);}
    private static List<String> stringList(RagQuery q,String key){Object value=q.getFilters().get(key);return value instanceof List<?> list?list.stream().map(String::valueOf).toList():List.of();}
    private static List<Map<String,Object>> maps(RagQuery q,String key){Object value=q.getFilters().get(key);if(!(value instanceof List<?> list))return List.of();return list.stream().filter(Map.class::isInstance).map(item->{Map<?,?> raw=(Map<?,?>)item;Map<String,Object> map=new LinkedHashMap<>();raw.forEach((k,v)->map.put(String.valueOf(k),v));return map;}).toList();}
    private static List<Map<String,Object>> normalizeExamples(List<Map<String,Object>> examples){return examples.stream().map(item->item.containsKey("cypher")?item:Map.<String,Object>of("question",item.get("question"),"cypher",item.get("query"))).toList();}
    private static Map<String,List<String>> properties(RagQuery q){Object value=q.getFilters().get("allowedProperties");if(!(value instanceof Map<?,?> map))return Map.of();Map<String,List<String>> result=new LinkedHashMap<>();map.forEach((key,raw)->{if(raw instanceof List<?> list)result.put(String.valueOf(key),list.stream().map(String::valueOf).toList());});return result;}
    private static Map<String,List<String>> allowedProperties(GraphSchemaSnapshot schema){Map<String,List<String>> result=new LinkedHashMap<>();schema.nodes().forEach(node->result.put(node.label(),new ArrayList<>(node.properties().keySet())));schema.relationships().forEach(rel->result.put(rel.type(),new ArrayList<>(rel.properties().keySet())));return result;}
}