diff --git a/deployment/config/java-shared/application.properties.sample b/deployment/config/java-shared/application.properties.sample index 765c71cd..3f794130 100644 --- a/deployment/config/java-shared/application.properties.sample +++ b/deployment/config/java-shared/application.properties.sample @@ -232,6 +232,7 @@ llm.sync.anthropic.api-key= #RAG codecrow.rag.api.url=http://host.docker.internal:8001 codecrow.rag.api.enabled=true +codecrow.rag.api.secret=change-me-to-a-random-secret # RAG API timeouts (in seconds) codecrow.rag.api.timeout.connect=30 codecrow.rag.api.timeout.read=120 diff --git a/deployment/config/mcp-client/.env.sample b/deployment/config/mcp-client/.env.sample index c92e89a9..43168803 100644 --- a/deployment/config/mcp-client/.env.sample +++ b/deployment/config/mcp-client/.env.sample @@ -2,6 +2,17 @@ AI_CLIENT_PORT=8000 RAG_ENABLED=true RAG_API_URL=http://host.docker.internal:8001 +# === Service-to-Service Auth === +# Shared secret for authenticating requests between internal services. +# Must match the SERVICE_SECRET configured on rag-pipeline. +# Leave empty to disable auth (dev mode only). +# IMPORTANT: Avoid $ { } characters in the secret — they can cause dotenv parsing issues. +SERVICE_SECRET=change-me-to-a-random-secret + +# === Concurrency === +# Max parallel review requests handled simultaneously (default: 4) +MAX_CONCURRENT_REVIEWS=4 + # API access for Platform MCP (internal network only) CODECROW_API_URL=http://codecrow-web-application:8081 diff --git a/deployment/config/rag-pipeline/.env.sample b/deployment/config/rag-pipeline/.env.sample index e19d8008..67003ea4 100644 --- a/deployment/config/rag-pipeline/.env.sample +++ b/deployment/config/rag-pipeline/.env.sample @@ -1,7 +1,37 @@ +# === Service-to-Service Auth === +# Shared secret for authenticating incoming requests from mcp-client. +# Must match the SERVICE_SECRET configured on mcp-client. +# Leave empty to disable auth (dev mode only). +# IMPORTANT: Avoid $ { } characters in the secret — they can cause dotenv parsing issues. +SERVICE_SECRET=change-me-to-a-random-secret + +# === Path Traversal Guard === +# Root directory that repo_path arguments are allowed under. +# The rag-pipeline will reject any index/query request whose resolved +# path escapes this directory. Default: /tmp +ALLOWED_REPO_ROOT=/tmp + #QDRANT configuration QDRANT_URL=http://qdrant:6333 QDRANT_COLLECTION_PREFIX=codecrow +# ollama/openrouter +EMBEDDING_PROVIDER=openrouter +OLLAMA_BASE_URL=http://localhost:11434 +OLLAMA_EMBEDDING_MODEL=qwen3-embedding:0.6b + +# Ollama Performance Tuning +# Batch size for embedding requests (higher = better throughput, more memory) +OLLAMA_BATCH_SIZE=100 +# Request timeout in seconds (increase for slow CPU) +OLLAMA_TIMEOUT=120 + +# CPU Threading Optimization (set based on your CPU cores) +# Recommended: physical_cores - 1 (leave 1 core for system) +OMP_NUM_THREADS=6 +MKL_NUM_THREADS=6 +OPENBLAS_NUM_THREADS=6 + # OpenRouter Configuration # Get your API key from https://openrouter.ai/ OPENROUTER_API_KEY=sk-or-v1-your-api-key-here diff --git a/deployment/config/web-frontend/.env.sample b/deployment/config/web-frontend/.env.sample index 40b36324..c9f9f269 100644 --- a/deployment/config/web-frontend/.env.sample +++ b/deployment/config/web-frontend/.env.sample @@ -3,6 +3,26 @@ VITE_WEBHOOK_URL=http://localhost:8082 VITE_GOOGLE_CLIENT_ID=YOUR_GOOGLE_CLIENT_ID SERVER_PORT=8080 +VITE_BLOG_URL=http://localhost:8083 + +# ============================================================================ +# Feature Flags (Cloud-specific features - disabled by default for OSS) +# ============================================================================ +# Set to "true" to enable cloud features. Leave empty or "false" to disable. +# +# VITE_FEATURE_BILLING - Enables billing page, subscription management, payment methods +VITE_FEATURE_BILLING=false +# +# VITE_FEATURE_CLOUD_PLANS - Enables cloud subscription plans (Pro, Pro+, Enterprise) +VITE_FEATURE_CLOUD_PLANS=false +# +# VITE_FEATURE_USAGE_ANALYTICS - Enables usage tracking and quota management +VITE_FEATURE_USAGE_ANALYTICS=false +# +# VITE_FEATURE_ENTERPRISE - Enables SSO, SAML, and advanced team management +VITE_FEATURE_ENTERPRISE=false + + # New Relic Browser Monitoring (Optional - leave empty to disable) # Get these values from: https://one.newrelic.com -> Browser -> Add data -> Browser monitoring # License Key: Use the BROWSER (Ingest) license key, NOT the main license key diff --git a/frontend b/frontend index 9ce813e0..d97b8264 160000 --- a/frontend +++ b/frontend @@ -1 +1 @@ -Subproject commit 9ce813e0ad5798c7059cf6c23c08b4d0a36c74b5 +Subproject commit d97b826464e13edb3da8a2b9a3e4b32680f23001 diff --git a/java-ecosystem/libs/analysis-engine/pom.xml b/java-ecosystem/libs/analysis-engine/pom.xml index 4d16d658..3fba0c89 100644 --- a/java-ecosystem/libs/analysis-engine/pom.xml +++ b/java-ecosystem/libs/analysis-engine/pom.xml @@ -68,6 +68,12 @@ okhttp + + + com.knuddels + jtokkit + + org.junit.jupiter diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/module-info.java b/java-ecosystem/libs/analysis-engine/src/main/java/module-info.java index b3e30345..c61e453f 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/module-info.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/module-info.java @@ -18,10 +18,12 @@ requires com.fasterxml.jackson.annotation; requires jakarta.persistence; requires kotlin.stdlib; + requires jtokkit; exports org.rostilos.codecrow.analysisengine.aiclient; exports org.rostilos.codecrow.analysisengine.config; exports org.rostilos.codecrow.analysisengine.dto.request.ai; + exports org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment; exports org.rostilos.codecrow.analysisengine.dto.request.processor; exports org.rostilos.codecrow.analysisengine.dto.request.validation; exports org.rostilos.codecrow.analysisengine.exception; diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/aiclient/AiAnalysisClient.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/aiclient/AiAnalysisClient.java index d31b136a..bf2b7a8b 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/aiclient/AiAnalysisClient.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/aiclient/AiAnalysisClient.java @@ -160,7 +160,7 @@ public Map performAnalysis(AiAnalysisRequest request, java.util. /** * Extracts analysis data from nested response structure. - * Expected: response -> result -> {comment, issues} + * Expected: response -> result -> {comment, issues, inference_stats} * Issues can be either a List (array) or Map (object with numeric keys) */ private Map extractAndValidateAnalysisData(Map response) throws IOException { @@ -176,6 +176,15 @@ private Map extractAndValidateAnalysisData(Map response) throws if (result == null) { throw new IOException("Missing 'result' field in AI response"); } + + // Check for error response from MCP client + Object errorFlag = result.get("error"); + if (Boolean.TRUE.equals(errorFlag) || "true".equals(String.valueOf(errorFlag))) { + String errorMessage = result.get("error_message") != null + ? String.valueOf(result.get("error_message")) + : String.valueOf(result.get("comment")); + throw new IOException("Analysis failed: " + errorMessage); + } if (!result.containsKey("comment") || !result.containsKey("issues")) { throw new IOException("Analysis data missing required fields: 'comment' and/or 'issues'"); diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiAnalysisRequestImpl.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiAnalysisRequestImpl.java index a8244eca..54f05b5a 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiAnalysisRequestImpl.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiAnalysisRequestImpl.java @@ -1,6 +1,7 @@ package org.rostilos.codecrow.analysisengine.dto.request.ai; import com.fasterxml.jackson.annotation.JsonProperty; +import org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment.PrEnrichmentDataDto; import org.rostilos.codecrow.core.model.ai.AIConnection; import org.rostilos.codecrow.core.model.ai.AIProviderKey; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisMode; @@ -44,6 +45,9 @@ public class AiAnalysisRequestImpl implements AiAnalysisRequest{ protected final String deltaDiff; protected final String previousCommitHash; protected final String currentCommitHash; + + // File enrichment data (full file contents + dependency graph) + protected final PrEnrichmentDataDto enrichmentData; protected AiAnalysisRequestImpl(Builder builder) { this.projectId = builder.projectId; @@ -74,6 +78,8 @@ protected AiAnalysisRequestImpl(Builder builder) { this.deltaDiff = builder.deltaDiff; this.previousCommitHash = builder.previousCommitHash; this.currentCommitHash = builder.currentCommitHash; + // File enrichment data + this.enrichmentData = builder.enrichmentData; } public Long getProjectId() { @@ -181,6 +187,10 @@ public String getCurrentCommitHash() { return currentCommitHash; } + public PrEnrichmentDataDto getEnrichmentData() { + return enrichmentData; + } + public static Builder builder() { return new Builder<>(); @@ -216,6 +226,8 @@ public static class Builder> { private String deltaDiff; private String previousCommitHash; private String currentCommitHash; + // File enrichment data + private PrEnrichmentDataDto enrichmentData; protected Builder() { } @@ -278,6 +290,113 @@ public T withPreviousAnalysisData(Optional optionalPreviousAnalysi return self(); } + /** + * Set previous issues from ALL PR analysis versions. + * This provides the LLM with complete issue history including resolved issues, + * helping it understand what was already found and fixed. + * + * Issues are deduplicated by fingerprint (file + line ±3 + severity + truncated reason). + * When duplicates exist across versions, we keep the most recent version's data + * but preserve resolved status if ANY version marked it resolved. + * + * @param allPrAnalyses List of all analyses for this PR, ordered by version DESC (newest first) + */ + public T withAllPrAnalysesData(List allPrAnalyses) { + if (allPrAnalyses == null || allPrAnalyses.isEmpty()) { + return self(); + } + + // Convert all issues to DTOs + List allIssues = allPrAnalyses.stream() + .flatMap(analysis -> analysis.getIssues().stream()) + .map(AiRequestPreviousIssueDTO::fromEntity) + .toList(); + + // Deduplicate: group by fingerprint, keep most recent version but preserve resolved status + java.util.Map deduped = new java.util.LinkedHashMap<>(); + + for (AiRequestPreviousIssueDTO issue : allIssues) { + String fingerprint = computeIssueFingerprint(issue); + AiRequestPreviousIssueDTO existing = deduped.get(fingerprint); + + if (existing == null) { + // First occurrence of this issue + deduped.put(fingerprint, issue); + } else { + // Duplicate found - keep the one with higher prVersion (more recent) + // But if older version is resolved and newer is not, preserve resolved status + int existingVersion = existing.prVersion() != null ? existing.prVersion() : 0; + int currentVersion = issue.prVersion() != null ? issue.prVersion() : 0; + + boolean existingResolved = "resolved".equalsIgnoreCase(existing.status()); + boolean currentResolved = "resolved".equalsIgnoreCase(issue.status()); + + if (currentVersion > existingVersion) { + // Current is newer - use it, but preserve resolved status if existing was resolved + if (existingResolved && !currentResolved) { + // Older version was resolved but newer one isn't marked - use resolved data from older + deduped.put(fingerprint, mergeResolvedStatus(issue, existing)); + } else { + deduped.put(fingerprint, issue); + } + } else if (existingVersion == currentVersion) { + // Same version - prefer resolved one + if (currentResolved && !existingResolved) { + deduped.put(fingerprint, issue); + } + } + // If existing is newer, keep it (already in map) + } + } + + this.previousCodeAnalysisIssues = new java.util.ArrayList<>(deduped.values()); + + return self(); + } + + /** + * Compute a fingerprint for an issue to detect duplicates across PR versions. + * Uses: file + normalized line (±3 tolerance) + severity + first 50 chars of reason. + */ + private String computeIssueFingerprint(AiRequestPreviousIssueDTO issue) { + String file = issue.file() != null ? issue.file() : ""; + // Normalize line to nearest multiple of 3 for tolerance + int lineGroup = issue.line() != null ? (issue.line() / 3) : 0; + String severity = issue.severity() != null ? issue.severity() : ""; + String reasonPrefix = issue.reason() != null + ? issue.reason().substring(0, Math.min(50, issue.reason().length())).toLowerCase().trim() + : ""; + + return file + "::" + lineGroup + "::" + severity + "::" + reasonPrefix; + } + + /** + * Merge resolved status from an older issue version into a newer one. + * Creates a new DTO with the newer issue's data but the older issue's resolution info. + */ + private AiRequestPreviousIssueDTO mergeResolvedStatus( + AiRequestPreviousIssueDTO newer, + AiRequestPreviousIssueDTO resolvedOlder) { + return new AiRequestPreviousIssueDTO( + newer.id(), + newer.type(), + newer.severity(), + newer.reason(), + newer.suggestedFixDescription(), + newer.suggestedFixDiff(), + newer.file(), + newer.line(), + newer.branch(), + newer.pullRequestId(), + resolvedOlder.status(), // Use resolved status from older + newer.category(), + newer.prVersion(), + resolvedOlder.resolvedDescription(), + resolvedOlder.resolvedByCommit(), + resolvedOlder.resolvedInAnalysisId() + ); + } + public T withMaxAllowedTokens(int maxAllowedTokens) { this.maxAllowedTokens = maxAllowedTokens; return self(); @@ -354,6 +473,11 @@ public T withCurrentCommitHash(String currentCommitHash) { return self(); } + public T withEnrichmentData(PrEnrichmentDataDto enrichmentData) { + this.enrichmentData = enrichmentData; + return self(); + } + public AiAnalysisRequestImpl build() { return new AiAnalysisRequestImpl(this); } diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTO.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTO.java index 8d166455..f14936c0 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTO.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTO.java @@ -15,12 +15,23 @@ public record AiRequestPreviousIssueDTO( String branch, String pullRequestId, String status, // open|resolved|ignored - String category + String category, + // Resolution tracking fields + Integer prVersion, // Which PR iteration this issue was found in + String resolvedDescription, // Description of how the issue was resolved + String resolvedByCommit, // Commit hash that resolved the issue + Long resolvedInAnalysisId // Analysis ID where this was resolved (null if still open) ) { public static AiRequestPreviousIssueDTO fromEntity(CodeAnalysisIssue issue) { String categoryStr = issue.getIssueCategory() != null ? issue.getIssueCategory().name() : IssueCategory.CODE_QUALITY.name(); + + Integer prVersion = null; + if (issue.getAnalysis() != null) { + prVersion = issue.getAnalysis().getPrVersion(); + } + return new AiRequestPreviousIssueDTO( String.valueOf(issue.getId()), categoryStr, @@ -33,7 +44,11 @@ public static AiRequestPreviousIssueDTO fromEntity(CodeAnalysisIssue issue) { issue.getAnalysis() == null ? null : issue.getAnalysis().getBranchName(), issue.getAnalysis() == null || issue.getAnalysis().getPrNumber() == null ? null : String.valueOf(issue.getAnalysis().getPrNumber()), issue.isResolved() ? "resolved" : "open", - categoryStr + categoryStr, + prVersion, + issue.getResolvedDescription(), + issue.getResolvedCommitHash(), + issue.getResolvedAnalysisId() ); } } diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/FileContentDto.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/FileContentDto.java new file mode 100644 index 00000000..85566f95 --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/FileContentDto.java @@ -0,0 +1,46 @@ +package org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment; + +/** + * DTO representing the content of a single file retrieved from VCS. + * Used for file enrichment during PR analysis to provide full file context. + */ +public record FileContentDto( + String path, + String content, + long sizeBytes, + boolean skipped, + String skipReason +) { + /** + * Create a successful file content result. + */ + public static FileContentDto of(String path, String content) { + return new FileContentDto( + path, + content, + content != null ? content.getBytes().length : 0, + false, + null + ); + } + + /** + * Create a skipped file result (e.g., file too large, binary, or fetch failed). + */ + public static FileContentDto skipped(String path, String reason) { + return new FileContentDto(path, null, 0, true, reason); + } + + /** + * Create a skipped file result due to size limit. + */ + public static FileContentDto skippedDueToSize(String path, long actualSize, long maxSize) { + return new FileContentDto( + path, + null, + actualSize, + true, + String.format("File size %d bytes exceeds limit %d bytes", actualSize, maxSize) + ); + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/FileRelationshipDto.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/FileRelationshipDto.java new file mode 100644 index 00000000..21adaa2d --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/FileRelationshipDto.java @@ -0,0 +1,90 @@ +package org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment; + +/** + * DTO representing a relationship between two files in the PR. + * Used for building the dependency graph for intelligent batching. + */ +public record FileRelationshipDto( + String sourceFile, + String targetFile, + RelationshipType relationshipType, + String matchedOn, + int strength +) { + /** + * Types of relationships between files. + */ + public enum RelationshipType { + IMPORTS, + EXTENDS, + IMPLEMENTS, + CALLS, + SAME_PACKAGE, + REFERENCES + } + + /** + * Create an import relationship. + */ + public static FileRelationshipDto imports(String sourceFile, String targetFile, String importStatement) { + return new FileRelationshipDto( + sourceFile, + targetFile, + RelationshipType.IMPORTS, + importStatement, + 10 // High strength for direct imports + ); + } + + /** + * Create an extends relationship. + */ + public static FileRelationshipDto extendsClass(String sourceFile, String targetFile, String className) { + return new FileRelationshipDto( + sourceFile, + targetFile, + RelationshipType.EXTENDS, + className, + 15 // Highest strength for inheritance + ); + } + + /** + * Create an implements relationship. + */ + public static FileRelationshipDto implementsInterface(String sourceFile, String targetFile, String interfaceName) { + return new FileRelationshipDto( + sourceFile, + targetFile, + RelationshipType.IMPLEMENTS, + interfaceName, + 15 // Highest strength for interface implementation + ); + } + + /** + * Create a calls relationship. + */ + public static FileRelationshipDto calls(String sourceFile, String targetFile, String methodName) { + return new FileRelationshipDto( + sourceFile, + targetFile, + RelationshipType.CALLS, + methodName, + 8 // Medium-high strength for method calls + ); + } + + /** + * Create a same-package relationship. + */ + public static FileRelationshipDto samePackage(String sourceFile, String targetFile, String packageName) { + return new FileRelationshipDto( + sourceFile, + targetFile, + RelationshipType.SAME_PACKAGE, + packageName, + 3 // Low strength for implicit package relationship + ); + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/ParsedFileMetadataDto.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/ParsedFileMetadataDto.java new file mode 100644 index 00000000..c6017ab5 --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/ParsedFileMetadataDto.java @@ -0,0 +1,70 @@ +package org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.List; + +/** + * DTO representing parsed AST metadata for a single file. + * Mirrors the response from RAG pipeline's /parse endpoint. + */ +@JsonIgnoreProperties(ignoreUnknown = true) +public record ParsedFileMetadataDto( + @JsonProperty("path") String path, + @JsonProperty("language") String language, + @JsonProperty("imports") List imports, + @JsonProperty("extends") List extendsClasses, + @JsonProperty("implements") List implementsInterfaces, + @JsonProperty("semantic_names") List semanticNames, + @JsonProperty("parent_class") String parentClass, + @JsonProperty("namespace") String namespace, + @JsonProperty("calls") List calls, + @JsonProperty("error") String error +) { + /** + * Create a metadata result with only imports and extends (minimal parsing). + */ + public static ParsedFileMetadataDto minimal(String path, List imports, List extendsClasses) { + return new ParsedFileMetadataDto( + path, + null, + imports, + extendsClasses, + List.of(), + List.of(), + null, + null, + List.of(), + null + ); + } + + /** + * Create an error result for a file that couldn't be parsed. + */ + public static ParsedFileMetadataDto error(String path, String errorMessage) { + return new ParsedFileMetadataDto( + path, + null, + List.of(), + List.of(), + List.of(), + List.of(), + null, + null, + List.of(), + errorMessage + ); + } + + /** + * Check if this metadata has any relationships to extract. + */ + public boolean hasRelationships() { + return (imports != null && !imports.isEmpty()) || + (extendsClasses != null && !extendsClasses.isEmpty()) || + (implementsInterfaces != null && !implementsInterfaces.isEmpty()) || + (calls != null && !calls.isEmpty()); + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/PrEnrichmentDataDto.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/PrEnrichmentDataDto.java new file mode 100644 index 00000000..2b1dec3b --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/ai/enrichment/PrEnrichmentDataDto.java @@ -0,0 +1,63 @@ +package org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment; + +import java.util.List; +import java.util.Map; + +/** + * Aggregate DTO containing all file enrichment data for a PR. + * This is the result of the enrichment process and can be serialized for transfer. + */ +public record PrEnrichmentDataDto( + List fileContents, + List fileMetadata, + List relationships, + EnrichmentStats stats +) { + /** + * Statistics about the enrichment process. + */ + public record EnrichmentStats( + int totalFilesRequested, + int filesEnriched, + int filesSkipped, + int relationshipsFound, + long totalContentSizeBytes, + long processingTimeMs, + Map skipReasons + ) { + public static EnrichmentStats empty() { + return new EnrichmentStats(0, 0, 0, 0, 0, 0, Map.of()); + } + } + + /** + * Create empty enrichment data (when enrichment is disabled or not applicable). + */ + public static PrEnrichmentDataDto empty() { + return new PrEnrichmentDataDto( + List.of(), + List.of(), + List.of(), + EnrichmentStats.empty() + ); + } + + /** + * Check if enrichment data is present. + */ + public boolean hasData() { + return (fileContents != null && !fileContents.isEmpty()) || + (relationships != null && !relationships.isEmpty()); + } + + /** + * Get total size of all file contents in bytes. + */ + public long getTotalContentSize() { + if (fileContents == null) return 0; + return fileContents.stream() + .filter(f -> !f.skipped()) + .mapToLong(FileContentDto::sizeBytes) + .sum(); + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/processor/PrProcessRequest.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/processor/PrProcessRequest.java index 5dc47f1e..752efc2a 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/processor/PrProcessRequest.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/dto/request/processor/PrProcessRequest.java @@ -35,6 +35,13 @@ public class PrProcessRequest implements AnalysisProcessRequest { public String prAuthorId; public String prAuthorUsername; + + /** + * Optional pre-acquired lock key. If set, the processor will skip lock acquisition + * and use this lock key directly. This prevents double-locking when the webhook handler + * has already acquired the lock before calling the processor. + */ + public String preAcquiredLockKey; public Long getProjectId() { @@ -64,4 +71,6 @@ public String getSourceBranchName() { public String getPrAuthorId() { return prAuthorId; } public String getPrAuthorUsername() { return prAuthorUsername; } + + public String getPreAcquiredLockKey() { return preAcquiredLockKey; } } diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/exception/DiffTooLargeException.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/exception/DiffTooLargeException.java new file mode 100644 index 00000000..7304448c --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/exception/DiffTooLargeException.java @@ -0,0 +1,47 @@ +package org.rostilos.codecrow.analysisengine.exception; + +/** + * Exception thrown when a diff exceeds the configured token limit for analysis. + * This is a soft skip - the analysis is not performed but the job is not marked as failed. + */ +public class DiffTooLargeException extends RuntimeException { + + private final int estimatedTokens; + private final int maxAllowedTokens; + private final Long projectId; + private final Long pullRequestId; + + public DiffTooLargeException(int estimatedTokens, int maxAllowedTokens, Long projectId, Long pullRequestId) { + super(String.format( + "PR diff exceeds token limit: estimated %d tokens, max allowed %d tokens (project=%d, PR=%d)", + estimatedTokens, maxAllowedTokens, projectId, pullRequestId + )); + this.estimatedTokens = estimatedTokens; + this.maxAllowedTokens = maxAllowedTokens; + this.projectId = projectId; + this.pullRequestId = pullRequestId; + } + + public int getEstimatedTokens() { + return estimatedTokens; + } + + public int getMaxAllowedTokens() { + return maxAllowedTokens; + } + + public Long getProjectId() { + return projectId; + } + + public Long getPullRequestId() { + return pullRequestId; + } + + /** + * Returns the percentage of the token limit that would be used. + */ + public double getUtilizationPercentage() { + return maxAllowedTokens > 0 ? (estimatedTokens * 100.0 / maxAllowedTokens) : 0; + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessor.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessor.java index 80f54fd0..41979a48 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessor.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessor.java @@ -224,10 +224,10 @@ public Map process(BranchProcessRequest request, Consumer existingFiles = updateBranchFiles(changedFiles, project, request.getTargetBranchName()); Branch branch = createOrUpdateProjectBranch(project, request); - mapCodeAnalysisIssuesToBranch(changedFiles, branch, project); + mapCodeAnalysisIssuesToBranch(changedFiles, existingFiles, branch, project); // Always update branch issue counts after mapping (even on first analysis) // Previously this was only done in reanalyzeCandidateIssues() which could be skipped @@ -282,10 +282,15 @@ public Set parseFilePathsFromDiff(String rawDiff) { return files; } - private void updateBranchFiles(Set changedFiles, Project project, String branchName) { + /** + * Updates branch file records for changed files. + * @return the set of file paths confirmed to exist in the branch (used to avoid redundant API calls) + */ + private Set updateBranchFiles(Set changedFiles, Project project, String branchName) { VcsInfo vcsInfo = getVcsInfo(project); EVcsProvider provider = getVcsProvider(project); VcsOperationsService operationsService = vcsServiceFactory.getOperationsService(provider); + Set filesExistingInBranch = new HashSet<>(); for (String filePath : changedFiles) { try { @@ -303,9 +308,12 @@ private void updateBranchFiles(Set changedFiles, Project project, String log.debug("Skipping file {} - does not exist in branch {}", filePath, branchName); continue; } + filesExistingInBranch.add(filePath); } catch (Exception e) { log.warn("Failed to check file existence for {} in branch {}: {}. Proceeding anyway.", filePath, branchName, e.getMessage()); + // On error, assume the file exists so we don't skip it + filesExistingInBranch.add(filePath); } List relatedIssues = codeAnalysisIssueRepository @@ -314,7 +322,14 @@ private void updateBranchFiles(Set changedFiles, Project project, String .filter(issue -> branchName.equals(issue.getAnalysis().getBranchName()) || branchName.equals(issue.getAnalysis().getSourceBranchName())) .toList(); - long unresolvedCount = branchSpecific.stream().filter(i -> !i.isResolved()).count(); + + // Deduplicate by content key before counting — multiple analyses may + // report the same logical issue with different DB ids + Set seenKeys = new HashSet<>(); + long unresolvedCount = branchSpecific.stream() + .filter(i -> !i.isResolved()) + .filter(i -> seenKeys.add(buildIssueContentKey(i))) + .count(); Optional projectFileOptional = branchFileRepository .findByProjectIdAndBranchNameAndFilePath(project.getId(), branchName, filePath); @@ -333,6 +348,7 @@ private void updateBranchFiles(Set changedFiles, Project project, String branchFileRepository.save(branchFile); } } + return filesExistingInBranch; } private Branch createOrUpdateProjectBranch(Project project, BranchProcessRequest request) { @@ -348,31 +364,14 @@ private Branch createOrUpdateProjectBranch(Project project, BranchProcessRequest return branchRepository.save(branch); } - private void mapCodeAnalysisIssuesToBranch(Set changedFiles, Branch branch, Project project) { - VcsInfo vcsInfo = getVcsInfo(project); - EVcsProvider provider = getVcsProvider(project); - VcsOperationsService operationsService = vcsServiceFactory.getOperationsService(provider); - + private void mapCodeAnalysisIssuesToBranch(Set changedFiles, Set filesExistingInBranch, + Branch branch, Project project) { for (String filePath : changedFiles) { - try { - OkHttpClient client = vcsClientProvider.getHttpClient(vcsInfo.vcsConnection()); - - boolean fileExistsInBranch = operationsService.checkFileExistsInBranch( - client, - vcsInfo.workspace(), - vcsInfo.repoSlug(), - branch.getBranchName(), - filePath - ); - - if (!fileExistsInBranch) { - log.debug("Skipping issue mapping for file {} - does not exist in branch {}", + // Use cached file existence from updateBranchFiles to avoid redundant API calls + if (!filesExistingInBranch.contains(filePath)) { + log.debug("Skipping issue mapping for file {} - does not exist in branch {} (cached)", filePath, branch.getBranchName()); - continue; - } - } catch (Exception e) { - log.warn("Failed to check file existence for {} in branch {}: {}. Proceeding with mapping.", - filePath, branch.getBranchName(), e.getMessage()); + continue; } List allIssues = codeAnalysisIssueRepository.findByProjectIdAndFilePath(project.getId(), filePath); @@ -387,27 +386,67 @@ private void mapCodeAnalysisIssuesToBranch(Set changedFiles, Branch bran }) .toList(); + // Content-based deduplication: build a map of existing BranchIssues by content key + // to prevent the same logical issue from being linked multiple times across analyses. + // Key = "lineNumber:severity:category" — unique enough within a single file context. + List existingBranchIssues = branchIssueRepository + .findUnresolvedByBranchIdAndFilePath(branch.getId(), filePath); + Map contentKeyMap = new HashMap<>(); + for (BranchIssue bi : existingBranchIssues) { + String key = buildIssueContentKey(bi.getCodeAnalysisIssue()); + contentKeyMap.putIfAbsent(key, bi); + } + + int skipped = 0; for (CodeAnalysisIssue issue : branchSpecificIssues) { + // Tier 1: exact ID match — same CodeAnalysisIssue already linked Optional existing = branchIssueRepository .findByBranchIdAndCodeAnalysisIssueId(branch.getId(), issue.getId()); - BranchIssue bc; + if (existing.isPresent()) { - bc = existing.get(); - bc.setSeverity(issue.getSeverity()); - branchIssueRepository.saveAndFlush(bc); - } else { - bc = new BranchIssue(); - bc.setBranch(branch); - bc.setCodeAnalysisIssue(issue); - bc.setResolved(issue.isResolved()); + BranchIssue bc = existing.get(); bc.setSeverity(issue.getSeverity()); - bc.setFirstDetectedPrNumber(issue.getAnalysis() != null ? issue.getAnalysis().getPrNumber() : null); branchIssueRepository.saveAndFlush(bc); + continue; } + + // Tier 2: content-based dedup — same logical issue from a different analysis + String contentKey = buildIssueContentKey(issue); + if (contentKeyMap.containsKey(contentKey)) { + skipped++; + continue; + } + + // No match — create new BranchIssue + BranchIssue bc = new BranchIssue(); + bc.setBranch(branch); + bc.setCodeAnalysisIssue(issue); + bc.setResolved(issue.isResolved()); + bc.setSeverity(issue.getSeverity()); + bc.setFirstDetectedPrNumber(issue.getAnalysis() != null ? issue.getAnalysis().getPrNumber() : null); + branchIssueRepository.saveAndFlush(bc); + // Register in map so subsequent issues in this batch also dedup + contentKeyMap.put(contentKey, bc); + } + + if (skipped > 0) { + log.debug("Skipped {} duplicate issue(s) for file {} in branch {}", + skipped, filePath, branch.getBranchName()); } } } + /** + * Builds a content key for deduplication of branch issues. + * Two CodeAnalysisIssue records with the same key represent the same logical issue. + */ + private String buildIssueContentKey(CodeAnalysisIssue issue) { + return issue.getFilePath() + ":" + + issue.getLineNumber() + ":" + + issue.getSeverity() + ":" + + issue.getIssueCategory(); + } + private void reanalyzeCandidateIssues(Set changedFiles, Branch branch, Project project, BranchProcessRequest request, Consumer> consumer) { List candidateBranchIssues = new ArrayList<>(); for (String filePath : changedFiles) { diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessor.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessor.java index 634f0419..052a8585 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessor.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessor.java @@ -29,9 +29,12 @@ import java.time.Duration; import java.time.Instant; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Optional; +import org.rostilos.codecrow.analysisengine.util.DiffFingerprintUtil; + /** * Generic service that handles pull request analysis. * Uses VCS-specific services via VcsServiceFactory for provider-specific operations. @@ -90,34 +93,45 @@ public Map process( // Publish analysis started event publishAnalysisStartedEvent(project, request, correlationId); - Optional lockKey = analysisLockService.acquireLockWithWait( - project, - request.getSourceBranchName(), - AnalysisLockType.PR_ANALYSIS, - request.getCommitHash(), - request.getPullRequestId(), - consumer::accept - ); - - if (lockKey.isEmpty()) { - String message = String.format( - "Failed to acquire lock after %d minutes for project=%s, PR=%d, branch=%s. Another analysis is still in progress.", - analysisLockService.getLockWaitTimeoutMinutes(), - project.getId(), - request.getPullRequestId(), - request.getSourceBranchName() - ); - log.warn(message); - - // Publish failed event due to lock timeout - publishAnalysisCompletedEvent(project, request, correlationId, startTime, - AnalysisCompletedEvent.CompletionStatus.FAILED, 0, 0, "Lock acquisition timeout"); - - throw new AnalysisLockedException( - AnalysisLockType.PR_ANALYSIS.name(), + // Check if a lock was already acquired by the caller (e.g., webhook handler) + // to prevent double-locking which causes unnecessary 2-minute waits + String lockKey; + boolean isPreAcquired = false; + if (request.getPreAcquiredLockKey() != null && !request.getPreAcquiredLockKey().isBlank()) { + lockKey = request.getPreAcquiredLockKey(); + isPreAcquired = true; + log.info("Using pre-acquired lock: {} for project={}, PR={}", lockKey, project.getId(), request.getPullRequestId()); + } else { + Optional acquiredLock = analysisLockService.acquireLockWithWait( + project, request.getSourceBranchName(), - project.getId() + AnalysisLockType.PR_ANALYSIS, + request.getCommitHash(), + request.getPullRequestId(), + consumer::accept ); + + if (acquiredLock.isEmpty()) { + String message = String.format( + "Failed to acquire lock after %d minutes for project=%s, PR=%d, branch=%s. Another analysis is still in progress.", + analysisLockService.getLockWaitTimeoutMinutes(), + project.getId(), + request.getPullRequestId(), + request.getSourceBranchName() + ); + log.warn(message); + + // Publish failed event due to lock timeout + publishAnalysisCompletedEvent(project, request, correlationId, startTime, + AnalysisCompletedEvent.CompletionStatus.FAILED, 0, 0, "Lock acquisition timeout"); + + throw new AnalysisLockedException( + AnalysisLockType.PR_ANALYSIS.name(), + request.getSourceBranchName(), + project.getId() + ); + } + lockKey = acquiredLock.get(); } try { @@ -139,16 +153,75 @@ public Map process( return Map.of("status", "cached", "cached", true); } - Optional previousAnalysis = codeAnalysisService.getPreviousVersionCodeAnalysis( + // --- Fallback cache: same commit hash, any PR number (handles close/reopen) --- + Optional commitHashHit = codeAnalysisService.getAnalysisByCommitHash( + project.getId(), request.getCommitHash()); + if (commitHashHit.isPresent()) { + log.info("Commit-hash cache hit for project={}, commit={} (source PR={}). Cloning for PR={}.", + project.getId(), request.getCommitHash(), + commitHashHit.get().getPrNumber(), request.getPullRequestId()); + CodeAnalysis cloned = codeAnalysisService.cloneAnalysisForPr( + commitHashHit.get(), project, request.getPullRequestId(), + request.getCommitHash(), request.getTargetBranchName(), + request.getSourceBranchName(), commitHashHit.get().getDiffFingerprint()); + try { + reportingService.postAnalysisResults(cloned, project, + request.getPullRequestId(), pullRequest.getId(), + request.getPlaceholderCommentId()); + } catch (IOException e) { + log.error("Failed to post commit-hash cached results to VCS: {}", e.getMessage(), e); + } + publishAnalysisCompletedEvent(project, request, correlationId, startTime, + AnalysisCompletedEvent.CompletionStatus.SUCCESS, 0, 0, null); + return Map.of("status", "cached_by_commit", "cached", true); + } + + // Get all previous analyses for this PR to provide full issue history to AI + List allPrAnalyses = codeAnalysisService.getAllPrAnalyses( project.getId(), request.getPullRequestId() ); + + // Get the most recent analysis for incremental diff calculation + Optional previousAnalysis = allPrAnalyses.isEmpty() + ? Optional.empty() + : Optional.of(allPrAnalyses.get(0)); - // Ensure branch index exists for target branch if configured + // Ensure branch index exists for TARGET branch (e.g., "1.2.1-rc") + // This is where the PR will merge TO - we want RAG context from this branch ensureRagIndexForTargetBranch(project, request.getTargetBranchName(), consumer); VcsAiClientService aiClientService = vcsServiceFactory.getAiClientService(provider); - AiAnalysisRequest aiRequest = aiClientService.buildAiAnalysisRequest(project, request, previousAnalysis); + AiAnalysisRequest aiRequest = aiClientService.buildAiAnalysisRequest( + project, request, previousAnalysis, allPrAnalyses); + + // --- Diff fingerprint cache: same code changes, different PR/commit --- + String diffFingerprint = DiffFingerprintUtil.compute(aiRequest.getRawDiff()); + if (diffFingerprint != null) { + Optional fingerprintHit = codeAnalysisService.getAnalysisByDiffFingerprint( + project.getId(), diffFingerprint); + if (fingerprintHit.isPresent()) { + log.info("Diff fingerprint cache hit for project={}, fingerprint={} (source PR={}). Cloning for PR={}.", + project.getId(), diffFingerprint.substring(0, 8) + "...", + fingerprintHit.get().getPrNumber(), request.getPullRequestId()); + // TODO: Option B — LIGHTWEIGHT mode: instead of full clone, reuse Stage 1 issues + // but re-run Stage 2 cross-file analysis against the new target branch context. + CodeAnalysis cloned = codeAnalysisService.cloneAnalysisForPr( + fingerprintHit.get(), project, request.getPullRequestId(), + request.getCommitHash(), request.getTargetBranchName(), + request.getSourceBranchName(), diffFingerprint); + try { + reportingService.postAnalysisResults(cloned, project, + request.getPullRequestId(), pullRequest.getId(), + request.getPlaceholderCommentId()); + } catch (IOException e) { + log.error("Failed to post fingerprint-cached results to VCS: {}", e.getMessage(), e); + } + publishAnalysisCompletedEvent(project, request, correlationId, startTime, + AnalysisCompletedEvent.CompletionStatus.SUCCESS, 0, 0, null); + return Map.of("status", "cached_by_fingerprint", "cached", true); + } + } Map aiResponse = aiAnalysisClient.performAnalysis(aiRequest, event -> { try { @@ -168,7 +241,8 @@ public Map process( request.getSourceBranchName(), request.getCommitHash(), request.getPrAuthorId(), - request.getPrAuthorUsername() + request.getPrAuthorUsername(), + diffFingerprint ); int issuesFound = newAnalysis.getIssues() != null ? newAnalysis.getIssues().size() : 0; @@ -208,7 +282,9 @@ public Map process( return Map.of("status", "error", "message", e.getMessage()); } finally { - analysisLockService.releaseLock(lockKey.get()); + if (!isPreAcquired) { + analysisLockService.releaseLock(lockKey); + } } } diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/PrFileEnrichmentService.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/PrFileEnrichmentService.java new file mode 100644 index 00000000..6df5ea55 --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/PrFileEnrichmentService.java @@ -0,0 +1,539 @@ +package org.rostilos.codecrow.analysisengine.service; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.ObjectMapper; +import okhttp3.*; +import org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment.*; +import org.rostilos.codecrow.vcsclient.VcsClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.stereotype.Service; + +import java.io.IOException; +import java.util.*; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +/** + * Service responsible for enriching PR analysis requests with full file contents + * and pre-computed dependency relationships. + * + * This service: + * 1. Fetches full file contents for changed files (with size limits) + * 2. Calls RAG pipeline's /parse endpoint to extract AST metadata + * 3. Builds a relationship graph from the parsed metadata + * 4. Returns enriched data for intelligent batching in Python + */ +@Service +public class PrFileEnrichmentService { + private static final Logger log = LoggerFactory.getLogger(PrFileEnrichmentService.class); + + private static final MediaType JSON_MEDIA_TYPE = MediaType.get("application/json; charset=utf-8"); + + @Value("${pr.enrichment.enabled:true}") + private boolean enrichmentEnabled; + + @Value("${pr.enrichment.max-file-size-bytes:102400}") // 100KB default + private long maxFileSizeBytes; + + @Value("${pr.enrichment.max-total-size-bytes:10485760}") // 10MB default + private long maxTotalSizeBytes; + + @Value("${pr.enrichment.rag-pipeline-url:http://localhost:8006}") + private String ragPipelineUrl; + + @Value("${pr.enrichment.request-timeout-seconds:60}") + private int requestTimeoutSeconds; + + private final ObjectMapper objectMapper; + private final OkHttpClient httpClient; + + public PrFileEnrichmentService() { + this.objectMapper = new ObjectMapper(); + this.httpClient = new OkHttpClient.Builder() + .connectTimeout(30, TimeUnit.SECONDS) + .readTimeout(60, TimeUnit.SECONDS) + .writeTimeout(30, TimeUnit.SECONDS) + .build(); + } + + /** + * Check if PR enrichment is enabled. + */ + public boolean isEnrichmentEnabled() { + return enrichmentEnabled; + } + + /** + * Enrich PR analysis with full file contents and dependency graph. + * + * @param vcsClient VCS client for fetching file contents + * @param workspace VCS workspace/owner + * @param repoSlug Repository slug + * @param branch Branch name or commit SHA for file content + * @param changedFiles List of changed file paths + * @return Enrichment data with file contents and relationships + */ + public PrEnrichmentDataDto enrichPrFiles( + VcsClient vcsClient, + String workspace, + String repoSlug, + String branch, + List changedFiles + ) { + if (!enrichmentEnabled) { + log.debug("PR enrichment is disabled"); + return PrEnrichmentDataDto.empty(); + } + + if (changedFiles == null || changedFiles.isEmpty()) { + log.debug("No changed files to enrich"); + return PrEnrichmentDataDto.empty(); + } + + long startTime = System.currentTimeMillis(); + Map skipReasons = new HashMap<>(); + + try { + // Step 1: Filter to supported file types + List supportedFiles = filterSupportedFiles(changedFiles, skipReasons); + if (supportedFiles.isEmpty()) { + log.info("No supported files to enrich after filtering"); + return createEmptyResultWithStats(changedFiles.size(), 0, skipReasons, startTime); + } + + // Step 2: Fetch file contents in batch + log.info("Fetching {} file contents for enrichment (branch: {})", supportedFiles.size(), branch); + Map fileContents = vcsClient.getFileContents( + workspace, repoSlug, supportedFiles, branch, (int) Math.min(maxFileSizeBytes, Integer.MAX_VALUE) + ); + + // Step 3: Build FileContentDto list + List contentDtos = buildFileContentDtos( + supportedFiles, fileContents, skipReasons + ); + + // Check total size limit + long totalSize = contentDtos.stream() + .filter(f -> !f.skipped()) + .mapToLong(FileContentDto::sizeBytes) + .sum(); + + if (totalSize > maxTotalSizeBytes) { + log.warn("Total file content size {} exceeds limit {} - truncating", + totalSize, maxTotalSizeBytes); + contentDtos = truncateToSizeLimit(contentDtos, maxTotalSizeBytes, skipReasons); + } + + // Step 4: Parse files to extract AST metadata + List metadata = parseFilesForMetadata(contentDtos); + + // Step 5: Build relationship graph from metadata + List relationships = buildRelationshipGraph( + metadata, changedFiles + ); + + long processingTime = System.currentTimeMillis() - startTime; + + // Build stats + int filesEnriched = (int) contentDtos.stream().filter(f -> !f.skipped()).count(); + int filesSkipped = changedFiles.size() - filesEnriched; + + PrEnrichmentDataDto.EnrichmentStats stats = new PrEnrichmentDataDto.EnrichmentStats( + changedFiles.size(), + filesEnriched, + filesSkipped, + relationships.size(), + totalSize, + processingTime, + skipReasons + ); + + log.info("PR enrichment completed: {} files enriched, {} skipped, {} relationships in {}ms", + filesEnriched, filesSkipped, relationships.size(), processingTime); + + return new PrEnrichmentDataDto(contentDtos, metadata, relationships, stats); + + } catch (Exception e) { + log.error("Failed to enrich PR files: {}", e.getMessage(), e); + return createEmptyResultWithStats( + changedFiles.size(), 0, + Map.of("error", changedFiles.size()), + startTime + ); + } + } + + /** + * Filter files to only those with supported extensions for parsing. + */ + private List filterSupportedFiles(List files, Map skipReasons) { + Set supportedExtensions = Set.of( + ".java", ".py", ".js", ".ts", ".jsx", ".tsx", + ".go", ".rs", ".rb", ".php", ".cs", ".cpp", ".c", ".h", + ".kt", ".scala", ".swift", ".m", ".mm" + ); + + List supported = new ArrayList<>(); + for (String file : files) { + String lower = file.toLowerCase(); + boolean isSupported = supportedExtensions.stream() + .anyMatch(lower::endsWith); + + if (isSupported) { + supported.add(file); + } else { + skipReasons.merge("unsupported_extension", 1, Integer::sum); + } + } + + return supported; + } + + /** + * Build FileContentDto list from fetched contents. + */ + private List buildFileContentDtos( + List requestedFiles, + Map fileContents, + Map skipReasons + ) { + List result = new ArrayList<>(); + + for (String path : requestedFiles) { + String content = fileContents.get(path); + + if (content == null) { + result.add(FileContentDto.skipped(path, "fetch_failed")); + skipReasons.merge("fetch_failed", 1, Integer::sum); + } else if (content.isEmpty()) { + result.add(FileContentDto.skipped(path, "empty_file")); + skipReasons.merge("empty_file", 1, Integer::sum); + } else { + result.add(FileContentDto.of(path, content)); + } + } + + return result; + } + + /** + * Truncate file list to stay within total size limit. + * Prioritizes smaller files to maximize coverage. + */ + private List truncateToSizeLimit( + List contents, + long maxTotalSize, + Map skipReasons + ) { + // Sort by size (smallest first) to include more files + List sorted = contents.stream() + .filter(f -> !f.skipped()) + .sorted(Comparator.comparingLong(FileContentDto::sizeBytes)) + .collect(Collectors.toCollection(ArrayList::new)); + + List result = new ArrayList<>(); + long currentSize = 0; + + for (FileContentDto file : sorted) { + if (currentSize + file.sizeBytes() <= maxTotalSize) { + result.add(file); + currentSize += file.sizeBytes(); + } else { + result.add(FileContentDto.skipped(file.path(), "total_size_limit_exceeded")); + skipReasons.merge("total_size_limit", 1, Integer::sum); + } + } + + // Add already-skipped files + contents.stream() + .filter(FileContentDto::skipped) + .forEach(result::add); + + return result; + } + + /** + * Call RAG pipeline's /parse/batch endpoint to extract AST metadata. + */ + private List parseFilesForMetadata(List contents) { + List filesToParse = contents.stream() + .filter(f -> !f.skipped()) + .toList(); + + if (filesToParse.isEmpty()) { + return Collections.emptyList(); + } + + try { + // Build batch request + List> files = filesToParse.stream() + .map(f -> Map.of("path", f.path(), "content", f.content())) + .toList(); + + Map requestBody = Map.of("files", files); + String jsonBody = objectMapper.writeValueAsString(requestBody); + + Request request = new Request.Builder() + .url(ragPipelineUrl + "/parse/batch") + .post(RequestBody.create(jsonBody, JSON_MEDIA_TYPE)) + .build(); + + try (Response response = httpClient.newCall(request).execute()) { + if (!response.isSuccessful()) { + log.warn("RAG /parse/batch failed with status {}: {}", + response.code(), response.message()); + return createFallbackMetadata(filesToParse); + } + + ResponseBody responseBody = response.body(); + if (responseBody == null) { + return createFallbackMetadata(filesToParse); + } + + ParseBatchResponse batchResponse = objectMapper.readValue( + responseBody.string(), + ParseBatchResponse.class + ); + + return batchResponse.results != null + ? batchResponse.results + : createFallbackMetadata(filesToParse); + } + + } catch (IOException e) { + log.warn("Failed to call RAG /parse/batch: {}", e.getMessage()); + return createFallbackMetadata(filesToParse); + } + } + + /** + * Create minimal metadata when RAG parsing fails. + */ + private List createFallbackMetadata(List files) { + return files.stream() + .map(f -> ParsedFileMetadataDto.error(f.path(), "parse_failed")) + .toList(); + } + + /** + * Build relationship graph from parsed metadata. + * Matches imports/extends/calls against changed files. + */ + private List buildRelationshipGraph( + List metadata, + List changedFiles + ) { + List relationships = new ArrayList<>(); + + // Build a map of possible file matches for quick lookup + Map nameToPath = buildNameToPathMap(changedFiles); + + for (ParsedFileMetadataDto file : metadata) { + if (file.error() != null) continue; + + // Process imports + if (file.imports() != null) { + for (String importStmt : file.imports()) { + String targetPath = findMatchingFile(importStmt, nameToPath, changedFiles); + if (targetPath != null && !targetPath.equals(file.path())) { + relationships.add(FileRelationshipDto.imports( + file.path(), targetPath, importStmt)); + } + } + } + + // Process extends + if (file.extendsClasses() != null) { + for (String className : file.extendsClasses()) { + String targetPath = findMatchingFile(className, nameToPath, changedFiles); + if (targetPath != null && !targetPath.equals(file.path())) { + relationships.add(FileRelationshipDto.extendsClass( + file.path(), targetPath, className)); + } + } + } + + // Process implements + if (file.implementsInterfaces() != null) { + for (String interfaceName : file.implementsInterfaces()) { + String targetPath = findMatchingFile(interfaceName, nameToPath, changedFiles); + if (targetPath != null && !targetPath.equals(file.path())) { + relationships.add(FileRelationshipDto.implementsInterface( + file.path(), targetPath, interfaceName)); + } + } + } + + // Process calls + if (file.calls() != null) { + for (String call : file.calls()) { + String targetPath = findMatchingFile(call, nameToPath, changedFiles); + if (targetPath != null && !targetPath.equals(file.path())) { + relationships.add(FileRelationshipDto.calls( + file.path(), targetPath, call)); + } + } + } + } + + // Add same-package relationships + addSamePackageRelationships(changedFiles, relationships); + + // Deduplicate relationships + return relationships.stream() + .distinct() + .toList(); + } + + /** + * Build a map of class/module names to file paths for matching. + */ + private Map buildNameToPathMap(List filePaths) { + Map map = new HashMap<>(); + + for (String path : filePaths) { + // Extract filename without extension + String fileName = path.contains("/") + ? path.substring(path.lastIndexOf('/') + 1) + : path; + String nameWithoutExt = fileName.contains(".") + ? fileName.substring(0, fileName.lastIndexOf('.')) + : fileName; + + map.put(nameWithoutExt, path); + map.put(nameWithoutExt.toLowerCase(), path); + + // For Java-style paths, also map the package structure + // e.g., com/example/MyClass.java -> MyClass + if (path.endsWith(".java")) { + String className = nameWithoutExt; + map.put(className, path); + + // Also map full package path without extension + String fullPath = path.replace('/', '.').replace(".java", ""); + map.put(fullPath, path); + } + + // For Python, map module paths + if (path.endsWith(".py")) { + String modulePath = path.replace('/', '.').replace(".py", ""); + map.put(modulePath, path); + } + } + + return map; + } + + /** + * Find a matching file path for an import/extends statement. + */ + private String findMatchingFile( + String reference, + Map nameToPath, + List changedFiles + ) { + if (reference == null || reference.isEmpty()) return null; + + // Try direct match + if (nameToPath.containsKey(reference)) { + return nameToPath.get(reference); + } + + // Extract last component (class name from qualified name) + String simpleName = reference.contains(".") + ? reference.substring(reference.lastIndexOf('.') + 1) + : reference; + + if (nameToPath.containsKey(simpleName)) { + return nameToPath.get(simpleName); + } + + // Try case-insensitive + String lowerName = simpleName.toLowerCase(); + if (nameToPath.containsKey(lowerName)) { + return nameToPath.get(lowerName); + } + + // Try partial path matching + String normalizedRef = reference.replace('.', '/'); + for (String path : changedFiles) { + if (path.contains(normalizedRef) || path.endsWith(simpleName + ".java") + || path.endsWith(simpleName + ".py") + || path.endsWith(simpleName + ".ts") + || path.endsWith(simpleName + ".js")) { + return path; + } + } + + return null; + } + + /** + * Add implicit relationships for files in the same package/directory. + */ + private void addSamePackageRelationships( + List changedFiles, + List relationships + ) { + // Group files by directory + Map> filesByDir = changedFiles.stream() + .collect(Collectors.groupingBy(path -> { + int lastSlash = path.lastIndexOf('/'); + return lastSlash > 0 ? path.substring(0, lastSlash) : ""; + })); + + // Add relationships within each directory + for (Map.Entry> entry : filesByDir.entrySet()) { + List filesInDir = entry.getValue(); + String packageName = entry.getKey(); + + if (filesInDir.size() > 1) { + for (int i = 0; i < filesInDir.size(); i++) { + for (int j = i + 1; j < filesInDir.size(); j++) { + relationships.add(FileRelationshipDto.samePackage( + filesInDir.get(i), filesInDir.get(j), packageName)); + } + } + } + } + } + + private PrEnrichmentDataDto createEmptyResultWithStats( + int totalFiles, + int enriched, + Map skipReasons, + long startTime + ) { + long processingTime = System.currentTimeMillis() - startTime; + return new PrEnrichmentDataDto( + Collections.emptyList(), + Collections.emptyList(), + Collections.emptyList(), + new PrEnrichmentDataDto.EnrichmentStats( + totalFiles, enriched, totalFiles - enriched, + 0, 0, processingTime, skipReasons + ) + ); + } + + /** + * Response DTO for /parse/batch endpoint. + */ + @JsonIgnoreProperties(ignoreUnknown = true) + private static class ParseBatchResponse { + @JsonProperty("results") + public List results; + + @JsonProperty("total_files") + public int totalFiles; + + @JsonProperty("successful") + public int successful; + + @JsonProperty("failed") + public int failed; + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/vcs/VcsAiClientService.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/vcs/VcsAiClientService.java index 1ee29fbd..3b894fa9 100644 --- a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/vcs/VcsAiClientService.java +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/service/vcs/VcsAiClientService.java @@ -7,6 +7,7 @@ import org.rostilos.codecrow.analysisengine.dto.request.processor.AnalysisProcessRequest; import java.security.GeneralSecurityException; +import java.util.List; import java.util.Optional; /** @@ -31,4 +32,23 @@ AiAnalysisRequest buildAiAnalysisRequest( AnalysisProcessRequest request, Optional previousAnalysis ) throws GeneralSecurityException; + + /** + * Builds an AI analysis request with full PR issue history. + * + * @param project The project being analyzed + * @param request The analysis process request + * @param previousAnalysis Optional previous analysis for incremental analysis (used for delta diff calculation) + * @param allPrAnalyses All analyses for this PR, ordered by version DESC (for issue history) + * @return The AI analysis request ready to be sent to the AI client + */ + default AiAnalysisRequest buildAiAnalysisRequest( + Project project, + AnalysisProcessRequest request, + Optional previousAnalysis, + List allPrAnalyses + ) throws GeneralSecurityException { + // Default implementation falls back to the previous method for backward compatibility + return buildAiAnalysisRequest(project, request, previousAnalysis); + } } diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/util/DiffFingerprintUtil.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/util/DiffFingerprintUtil.java new file mode 100644 index 00000000..6f8a381e --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/util/DiffFingerprintUtil.java @@ -0,0 +1,102 @@ +package org.rostilos.codecrow.analysisengine.util; + +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Computes a content-based fingerprint of a unified diff. + *

+ * Only actual change lines ({@code +} / {@code -}) are included — context lines, + * hunk headers ({@code @@}), and file headers ({@code +++} / {@code ---} / {@code diff --git}) + * are excluded. The change lines are sorted to make the fingerprint stable regardless + * of file ordering within the diff. + *

+ * This allows detecting that two PRs carry the same code changes even if they target + * different branches (different merge-base → different context/hunk headers). + */ +public final class DiffFingerprintUtil { + + private DiffFingerprintUtil() { /* utility */ } + + /** + * Compute a SHA-256 hex digest of the normalised change lines in the given diff. + * + * @param rawDiff the filtered unified diff (may be {@code null} or empty) + * @return 64-char lowercase hex string, or {@code null} if the diff is blank + */ + public static String compute(String rawDiff) { + if (rawDiff == null || rawDiff.isBlank()) { + return null; + } + + List changeLines = extractChangeLines(rawDiff); + if (changeLines.isEmpty()) { + return null; + } + + // Sort for stability across different file orderings + Collections.sort(changeLines); + + try { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + for (String line : changeLines) { + digest.update(line.getBytes(StandardCharsets.UTF_8)); + digest.update((byte) '\n'); + } + return bytesToHex(digest.digest()); + } catch (NoSuchAlgorithmException e) { + // SHA-256 is guaranteed by the JVM spec — should never happen + throw new IllegalStateException("SHA-256 not available", e); + } + } + + /** + * Extract only the actual change lines from a unified diff. + * A "change line" starts with exactly one {@code +} or {@code -} and is NOT + * a file header ({@code +++}, {@code ---}) or a diff metadata line. + */ + private static List extractChangeLines(String diff) { + List lines = new ArrayList<>(); + // Normalise line endings + String normalised = diff.replace("\r\n", "\n").replace("\r", "\n"); + for (String raw : normalised.split("\n")) { + String line = trimTrailingWhitespace(raw); + if (line.isEmpty()) { + continue; + } + char first = line.charAt(0); + if (first != '+' && first != '-') { + continue; + } + // Skip file-level headers: "+++", "---", "diff --git" + if (line.startsWith("+++") || line.startsWith("---")) { + continue; + } + if (line.startsWith("diff ")) { + continue; + } + lines.add(line); + } + return lines; + } + + private static String trimTrailingWhitespace(String s) { + int end = s.length(); + while (end > 0 && Character.isWhitespace(s.charAt(end - 1))) { + end--; + } + return s.substring(0, end); + } + + private static String bytesToHex(byte[] bytes) { + StringBuilder sb = new StringBuilder(bytes.length * 2); + for (byte b : bytes) { + sb.append(String.format("%02x", b & 0xff)); + } + return sb.toString(); + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/util/TokenEstimator.java b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/util/TokenEstimator.java new file mode 100644 index 00000000..86a7da1a --- /dev/null +++ b/java-ecosystem/libs/analysis-engine/src/main/java/org/rostilos/codecrow/analysisengine/util/TokenEstimator.java @@ -0,0 +1,88 @@ +package org.rostilos.codecrow.analysisengine.util; + +import com.knuddels.jtokkit.Encodings; +import com.knuddels.jtokkit.api.Encoding; +import com.knuddels.jtokkit.api.EncodingRegistry; +import com.knuddels.jtokkit.api.EncodingType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Utility class for estimating token counts in text content. + * Uses the cl100k_base encoding (used by GPT-4, Claude, and most modern LLMs). + */ +public final class TokenEstimator { + private static final Logger log = LoggerFactory.getLogger(TokenEstimator.class); + + private static final EncodingRegistry ENCODING_REGISTRY = Encodings.newDefaultEncodingRegistry(); + private static final Encoding ENCODING = ENCODING_REGISTRY.getEncoding(EncodingType.CL100K_BASE); + + // Prevent instantiation + private TokenEstimator() { + throw new UnsupportedOperationException("Utility class"); + } + + /** + * Estimate the number of tokens in the given text. + * + * @param text The text to estimate tokens for + * @return The estimated token count, or 0 if text is null/empty + */ + public static int estimateTokens(String text) { + if (text == null || text.isEmpty()) { + return 0; + } + try { + return ENCODING.countTokens(text); + } catch (Exception e) { + log.warn("Failed to count tokens, using fallback estimation: {}", e.getMessage()); + // Fallback: rough estimate of ~4 characters per token + return text.length() / 4; + } + } + + /** + * Check if the estimated token count exceeds the given limit. + * + * @param text The text to check + * @param maxTokens The maximum allowed tokens + * @return true if the text exceeds the limit, false otherwise + */ + public static boolean exceedsLimit(String text, int maxTokens) { + return estimateTokens(text) > maxTokens; + } + + /** + * Result of a token estimation check with details. + */ + public record TokenEstimationResult( + int estimatedTokens, + int maxAllowedTokens, + boolean exceedsLimit, + double utilizationPercentage + ) { + public String toLogString() { + return String.format("Tokens: %d / %d (%.1f%%) - %s", + estimatedTokens, maxAllowedTokens, utilizationPercentage, + exceedsLimit ? "EXCEEDS LIMIT" : "within limit"); + } + } + + /** + * Estimate tokens and check against limit, returning detailed result. + * + * @param text The text to check + * @param maxTokens The maximum allowed tokens + * @return Detailed estimation result + */ + public static TokenEstimationResult estimateAndCheck(String text, int maxTokens) { + int estimated = estimateTokens(text); + double utilization = maxTokens > 0 ? (estimated * 100.0 / maxTokens) : 0; + return new TokenEstimationResult( + estimated, + maxTokens, + estimated > maxTokens, + utilization + ); + } +} diff --git a/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTOTest.java b/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTOTest.java index 2fe1aca2..ebeb50a7 100644 --- a/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTOTest.java +++ b/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/dto/request/ai/AiRequestPreviousIssueDTOTest.java @@ -30,7 +30,11 @@ void shouldCreateRecordWithAllFields() { "main", "100", "open", - "SECURITY" + "SECURITY", + 1, // prVersion + null, // resolvedDescription + null, // resolvedByCommit + null // resolvedInPrVersion ); assertThat(dto.id()).isEqualTo("123"); @@ -45,13 +49,14 @@ void shouldCreateRecordWithAllFields() { assertThat(dto.pullRequestId()).isEqualTo("100"); assertThat(dto.status()).isEqualTo("open"); assertThat(dto.category()).isEqualTo("SECURITY"); + assertThat(dto.prVersion()).isEqualTo(1); } @Test @DisplayName("should handle null values") void shouldHandleNullValues() { AiRequestPreviousIssueDTO dto = new AiRequestPreviousIssueDTO( - null, null, null, null, null, null, null, null, null, null, null, null + null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null ); assertThat(dto.id()).isNull(); @@ -64,13 +69,13 @@ void shouldHandleNullValues() { @DisplayName("should implement equals correctly") void shouldImplementEqualsCorrectly() { AiRequestPreviousIssueDTO dto1 = new AiRequestPreviousIssueDTO( - "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat" + "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat", 1, null, null, null ); AiRequestPreviousIssueDTO dto2 = new AiRequestPreviousIssueDTO( - "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat" + "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat", 1, null, null, null ); AiRequestPreviousIssueDTO dto3 = new AiRequestPreviousIssueDTO( - "2", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat" + "2", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat", 1, null, null, null ); assertThat(dto1).isEqualTo(dto2); @@ -81,10 +86,10 @@ void shouldImplementEqualsCorrectly() { @DisplayName("should implement hashCode correctly") void shouldImplementHashCodeCorrectly() { AiRequestPreviousIssueDTO dto1 = new AiRequestPreviousIssueDTO( - "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat" + "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat", 1, null, null, null ); AiRequestPreviousIssueDTO dto2 = new AiRequestPreviousIssueDTO( - "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat" + "1", "type", "HIGH", "reason", "fix", "diff", "file.java", 10, "main", "1", "open", "cat", 1, null, null, null ); assertThat(dto1.hashCode()).isEqualTo(dto2.hashCode()); @@ -94,10 +99,14 @@ void shouldImplementHashCodeCorrectly() { @DisplayName("should support resolved status") void shouldSupportResolvedStatus() { AiRequestPreviousIssueDTO dto = new AiRequestPreviousIssueDTO( - "1", "type", "LOW", "reason", null, null, "file.java", 5, "dev", "2", "resolved", "CODE_QUALITY" + "1", "type", "LOW", "reason", null, null, "file.java", 5, "dev", "2", "resolved", "CODE_QUALITY", + 1, "Fixed by adding null check", "abc123", 2L ); assertThat(dto.status()).isEqualTo("resolved"); + assertThat(dto.resolvedDescription()).isEqualTo("Fixed by adding null check"); + assertThat(dto.resolvedByCommit()).isEqualTo("abc123"); + assertThat(dto.resolvedInAnalysisId()).isEqualTo(2L); } @Nested @@ -110,6 +119,7 @@ void shouldConvertEntityWithAllFields() { CodeAnalysis analysis = mock(CodeAnalysis.class); when(analysis.getBranchName()).thenReturn("feature-branch"); when(analysis.getPrNumber()).thenReturn(42L); + when(analysis.getPrVersion()).thenReturn(2); CodeAnalysisIssue issue = mock(CodeAnalysisIssue.class); when(issue.getId()).thenReturn(123L); @@ -122,6 +132,9 @@ void shouldConvertEntityWithAllFields() { when(issue.getFilePath()).thenReturn("src/Main.java"); when(issue.getLineNumber()).thenReturn(50); when(issue.isResolved()).thenReturn(false); + when(issue.getResolvedDescription()).thenReturn(null); + when(issue.getResolvedCommitHash()).thenReturn(null); + when(issue.getResolvedAnalysisId()).thenReturn(null); AiRequestPreviousIssueDTO dto = AiRequestPreviousIssueDTO.fromEntity(issue); @@ -137,14 +150,16 @@ void shouldConvertEntityWithAllFields() { assertThat(dto.pullRequestId()).isEqualTo("42"); assertThat(dto.status()).isEqualTo("open"); assertThat(dto.category()).isEqualTo("SECURITY"); + assertThat(dto.prVersion()).isEqualTo(2); } @Test - @DisplayName("should convert resolved entity") + @DisplayName("should convert resolved entity with resolution tracking") void shouldConvertResolvedEntity() { CodeAnalysis analysis = mock(CodeAnalysis.class); when(analysis.getBranchName()).thenReturn("main"); when(analysis.getPrNumber()).thenReturn(10L); + when(analysis.getPrVersion()).thenReturn(3); CodeAnalysisIssue issue = mock(CodeAnalysisIssue.class); when(issue.getId()).thenReturn(456L); @@ -155,10 +170,17 @@ void shouldConvertResolvedEntity() { when(issue.getFilePath()).thenReturn("src/Utils.java"); when(issue.getLineNumber()).thenReturn(10); when(issue.isResolved()).thenReturn(true); + when(issue.getResolvedDescription()).thenReturn("Fixed by refactoring"); + when(issue.getResolvedCommitHash()).thenReturn("abc123def"); + when(issue.getResolvedAnalysisId()).thenReturn(5L); AiRequestPreviousIssueDTO dto = AiRequestPreviousIssueDTO.fromEntity(issue); assertThat(dto.status()).isEqualTo("resolved"); + assertThat(dto.prVersion()).isEqualTo(3); + assertThat(dto.resolvedDescription()).isEqualTo("Fixed by refactoring"); + assertThat(dto.resolvedByCommit()).isEqualTo("abc123def"); + assertThat(dto.resolvedInAnalysisId()).isEqualTo(5L); } @Test @@ -167,6 +189,7 @@ void shouldHandleNullIssueCategoryWithDefault() { CodeAnalysis analysis = mock(CodeAnalysis.class); when(analysis.getBranchName()).thenReturn("main"); when(analysis.getPrNumber()).thenReturn(1L); + when(analysis.getPrVersion()).thenReturn(1); CodeAnalysisIssue issue = mock(CodeAnalysisIssue.class); when(issue.getId()).thenReturn(1L); @@ -186,6 +209,7 @@ void shouldHandleNullIssueCategoryWithDefault() { void shouldHandleNullSeverity() { CodeAnalysis analysis = mock(CodeAnalysis.class); when(analysis.getBranchName()).thenReturn("main"); + when(analysis.getPrVersion()).thenReturn(1); CodeAnalysisIssue issue = mock(CodeAnalysisIssue.class); when(issue.getId()).thenReturn(2L); @@ -213,6 +237,7 @@ void shouldHandleNullAnalysis() { assertThat(dto.branch()).isNull(); assertThat(dto.pullRequestId()).isNull(); + assertThat(dto.prVersion()).isNull(); } @Test @@ -221,6 +246,7 @@ void shouldHandleAnalysisWithNullPrNumber() { CodeAnalysis analysis = mock(CodeAnalysis.class); when(analysis.getBranchName()).thenReturn("develop"); when(analysis.getPrNumber()).thenReturn(null); + when(analysis.getPrVersion()).thenReturn(1); CodeAnalysisIssue issue = mock(CodeAnalysisIssue.class); when(issue.getId()).thenReturn(4L); diff --git a/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessorTest.java b/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessorTest.java index c38168d3..6e7d5bac 100644 --- a/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessorTest.java +++ b/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/BranchAnalysisProcessorTest.java @@ -28,7 +28,6 @@ import org.rostilos.codecrow.core.persistence.repository.branch.BranchRepository; import org.rostilos.codecrow.core.persistence.repository.codeanalysis.CodeAnalysisIssueRepository; import org.rostilos.codecrow.vcsclient.VcsClientProvider; -import org.springframework.context.ApplicationEventPublisher; import java.io.IOException; import java.util.*; @@ -73,9 +72,6 @@ class BranchAnalysisProcessorTest { @Mock private RagOperationsService ragOperationsService; - @Mock - private ApplicationEventPublisher eventPublisher; - @Mock private VcsOperationsService operationsService; @@ -255,14 +251,14 @@ void shouldThrowAnalysisLockedExceptionWhenLockCannotBeAcquired() throws IOExcep when(projectService.getProjectWithConnections(1L)).thenReturn(project); when(project.getId()).thenReturn(1L); - when(project.getName()).thenReturn("Test Project"); when(analysisLockService.acquireLockWithWait(any(), anyString(), any(), anyString(), any(), any())) .thenReturn(Optional.empty()); assertThatThrownBy(() -> processor.process(request, consumer)) .isInstanceOf(AnalysisLockedException.class); - verify(eventPublisher, times(2)).publishEvent(any()); + // No consumer or event interactions should occur when lock is not acquired + verifyNoInteractions(consumer); } // Note: Full process() integration tests are complex and require extensive mocking. diff --git a/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessorTest.java b/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessorTest.java index 0b36db83..0a8c555d 100644 --- a/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessorTest.java +++ b/java-ecosystem/libs/analysis-engine/src/test/java/org/rostilos/codecrow/analysisengine/processor/analysis/PullRequestAnalysisProcessorTest.java @@ -140,10 +140,13 @@ void shouldSuccessfullyProcessPRAnalysis() throws Exception { when(codeAnalysisService.getCodeAnalysisCache(anyLong(), anyString(), anyLong())) .thenReturn(Optional.empty()); - when(codeAnalysisService.getPreviousVersionCodeAnalysis(anyLong(), anyLong())) + when(codeAnalysisService.getAnalysisByCommitHash(anyLong(), anyString())) .thenReturn(Optional.empty()); + when(codeAnalysisService.getAllPrAnalyses(anyLong(), anyLong())) + .thenReturn(List.of()); - when(aiClientService.buildAiAnalysisRequest(any(), any(), any())).thenReturn(aiAnalysisRequest); + when(aiClientService.buildAiAnalysisRequest(any(), any(), any(), anyList())).thenReturn(aiAnalysisRequest); + when(aiAnalysisRequest.getRawDiff()).thenReturn(""); Map aiResponse = Map.of( "comment", "Review comment", @@ -151,7 +154,8 @@ void shouldSuccessfullyProcessPRAnalysis() throws Exception { ); when(aiAnalysisClient.performAnalysis(any(), any())).thenReturn(aiResponse); - when(codeAnalysisService.createAnalysisFromAiResponse(any(), any(), anyLong(), anyString(), anyString(), anyString(), any(), any())) + when(codeAnalysisService.createAnalysisFromAiResponse( + any(), any(), anyLong(), anyString(), anyString(), anyString(), any(), any(), any())) .thenReturn(codeAnalysis); Map result = processor.process(request, consumer, project); diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTO.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTO.java index b9e16fb5..b04b4434 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTO.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTO.java @@ -11,8 +11,7 @@ public record AIConnectionDTO( AIProviderKey providerKey, String aiModel, OffsetDateTime createdAt, - OffsetDateTime updatedAt, - int tokenLimitation + OffsetDateTime updatedAt ) { public static AIConnectionDTO fromAiConnection(AIConnection aiConnection) { @@ -22,8 +21,7 @@ public static AIConnectionDTO fromAiConnection(AIConnection aiConnection) { aiConnection.getProviderKey(), aiConnection.getAiModel(), aiConnection.getCreatedAt(), - aiConnection.getUpdatedAt(), - aiConnection.getTokenLimitation() + aiConnection.getUpdatedAt() ); } } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/project/ProjectDTO.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/project/ProjectDTO.java index 9c2a70a3..98027a0d 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/project/ProjectDTO.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/dto/project/ProjectDTO.java @@ -32,7 +32,8 @@ public record ProjectDTO( String installationMethod, CommentCommandsConfigDTO commentCommandsConfig, Boolean webhooksConfigured, - Long qualityGateId + Long qualityGateId, + Integer maxAnalysisTokenLimit ) { public static ProjectDTO fromProject(Project project) { Long vcsConnectionId = null; @@ -123,6 +124,9 @@ public static ProjectDTO fromProject(Project project) { if (project.getVcsRepoBinding() != null) { webhooksConfigured = project.getVcsRepoBinding().isWebhooksConfigured(); } + + // Get maxAnalysisTokenLimit from config + Integer maxAnalysisTokenLimit = config != null ? config.maxAnalysisTokenLimit() : ProjectConfig.DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT; return new ProjectDTO( project.getId(), @@ -146,7 +150,8 @@ public static ProjectDTO fromProject(Project project) { installationMethod, commentCommandsConfigDTO, webhooksConfigured, - project.getQualityGate() != null ? project.getQualityGate().getId() : null + project.getQualityGate() != null ? project.getQualityGate().getId() : null, + maxAnalysisTokenLimit ); } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/ai/AIConnection.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/ai/AIConnection.java index 2ca682c3..f6558f75 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/ai/AIConnection.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/ai/AIConnection.java @@ -39,9 +39,6 @@ public class AIConnection { @Column(name = "updated_at", nullable = false) private OffsetDateTime updatedAt = OffsetDateTime.now(); - @Column(name= "token_limitation", nullable = false) - private int tokenLimitation = 100000; - @PreUpdate public void onUpdate() { this.updatedAt = OffsetDateTime.now(); @@ -98,12 +95,4 @@ public OffsetDateTime getCreatedAt() { public OffsetDateTime getUpdatedAt() { return updatedAt; } - - public void setTokenLimitation(int tokenLimitation) { - this.tokenLimitation = tokenLimitation; - } - - public int getTokenLimitation() { - return tokenLimitation; - } } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/codeanalysis/CodeAnalysis.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/codeanalysis/CodeAnalysis.java index ae4377fa..14bd145d 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/codeanalysis/CodeAnalysis.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/codeanalysis/CodeAnalysis.java @@ -38,6 +38,9 @@ public class CodeAnalysis { @Column(name = "commit_hash", length = 40) private String commitHash; + @Column(name = "diff_fingerprint", length = 64) + private String diffFingerprint; + @Column(name = "target_branch_name") private String branchName; @@ -113,6 +116,9 @@ public void updateIssueCounts() { public String getCommitHash() { return commitHash; } public void setCommitHash(String commitHash) { this.commitHash = commitHash; } + public String getDiffFingerprint() { return diffFingerprint; } + public void setDiffFingerprint(String diffFingerprint) { this.diffFingerprint = diffFingerprint; } + public String getBranchName() { return branchName; } public void setBranchName(String branchName) { this.branchName = branchName; } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/Project.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/Project.java index 1a956edd..b12b4227 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/Project.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/Project.java @@ -222,6 +222,15 @@ public void setConfiguration(org.rostilos.codecrow.core.model.project.config.Pro this.configuration = configuration; } + /** + * Returns the effective project configuration. + * If configuration is null, returns a new default ProjectConfig. + * This ensures callers always get a valid config with default values. + */ + public org.rostilos.codecrow.core.model.project.config.ProjectConfig getEffectiveConfig() { + return configuration != null ? configuration : new org.rostilos.codecrow.core.model.project.config.ProjectConfig(); + } + public org.rostilos.codecrow.core.model.branch.Branch getDefaultBranch() { return defaultBranch; } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/config/ProjectConfig.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/config/ProjectConfig.java index 66335185..99d18764 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/config/ProjectConfig.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/project/config/ProjectConfig.java @@ -24,6 +24,8 @@ * - branchAnalysisEnabled: whether to analyze branch pushes (default: true). * - installationMethod: how the project integration is installed (WEBHOOK, PIPELINE, GITHUB_ACTION). * - commentCommands: configuration for PR comment-triggered commands (/codecrow analyze, summarize, ask). + * - maxAnalysisTokenLimit: maximum allowed tokens for PR analysis (default: 200000). + * Analysis will be skipped if the diff exceeds this limit. * * @see BranchAnalysisConfig * @see RagConfig @@ -32,6 +34,8 @@ */ @JsonIgnoreProperties(ignoreUnknown = true) public class ProjectConfig { + public static final int DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT = 200000; + @JsonProperty("useLocalMcp") private boolean useLocalMcp; @@ -56,16 +60,27 @@ public class ProjectConfig { private InstallationMethod installationMethod; @JsonProperty("commentCommands") private CommentCommandsConfig commentCommands; + @JsonProperty("maxAnalysisTokenLimit") + private Integer maxAnalysisTokenLimit; public ProjectConfig() { this.useLocalMcp = false; this.prAnalysisEnabled = true; this.branchAnalysisEnabled = true; + this.maxAnalysisTokenLimit = DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT; } public ProjectConfig(boolean useLocalMcp, String mainBranch, BranchAnalysisConfig branchAnalysis, RagConfig ragConfig, Boolean prAnalysisEnabled, Boolean branchAnalysisEnabled, InstallationMethod installationMethod, CommentCommandsConfig commentCommands) { + this(useLocalMcp, mainBranch, branchAnalysis, ragConfig, prAnalysisEnabled, branchAnalysisEnabled, + installationMethod, commentCommands, DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT); + } + + public ProjectConfig(boolean useLocalMcp, String mainBranch, BranchAnalysisConfig branchAnalysis, + RagConfig ragConfig, Boolean prAnalysisEnabled, Boolean branchAnalysisEnabled, + InstallationMethod installationMethod, CommentCommandsConfig commentCommands, + Integer maxAnalysisTokenLimit) { this.useLocalMcp = useLocalMcp; this.mainBranch = mainBranch; this.defaultBranch = mainBranch; // Keep in sync for backward compatibility @@ -75,6 +90,7 @@ public ProjectConfig(boolean useLocalMcp, String mainBranch, BranchAnalysisConfi this.branchAnalysisEnabled = branchAnalysisEnabled; this.installationMethod = installationMethod; this.commentCommands = commentCommands; + this.maxAnalysisTokenLimit = maxAnalysisTokenLimit != null ? maxAnalysisTokenLimit : DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT; } public ProjectConfig(boolean useLocalMcp, String mainBranch) { @@ -112,6 +128,14 @@ public String defaultBranch() { public InstallationMethod installationMethod() { return installationMethod; } public CommentCommandsConfig commentCommands() { return commentCommands; } + /** + * Get the maximum token limit for PR analysis. + * Returns the configured value or the default (200000) if not set. + */ + public int maxAnalysisTokenLimit() { + return maxAnalysisTokenLimit != null ? maxAnalysisTokenLimit : DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT; + } + // Setters for Jackson public void setUseLocalMcp(boolean useLocalMcp) { this.useLocalMcp = useLocalMcp; } @@ -149,6 +173,9 @@ public void setDefaultBranch(String defaultBranch) { public void setBranchAnalysisEnabled(Boolean branchAnalysisEnabled) { this.branchAnalysisEnabled = branchAnalysisEnabled; } public void setInstallationMethod(InstallationMethod installationMethod) { this.installationMethod = installationMethod; } public void setCommentCommands(CommentCommandsConfig commentCommands) { this.commentCommands = commentCommands; } + public void setMaxAnalysisTokenLimit(Integer maxAnalysisTokenLimit) { + this.maxAnalysisTokenLimit = maxAnalysisTokenLimit != null ? maxAnalysisTokenLimit : DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT; + } public void ensureMainBranchInPatterns() { String main = mainBranch(); @@ -230,13 +257,15 @@ public boolean equals(Object o) { Objects.equals(prAnalysisEnabled, that.prAnalysisEnabled) && Objects.equals(branchAnalysisEnabled, that.branchAnalysisEnabled) && installationMethod == that.installationMethod && - Objects.equals(commentCommands, that.commentCommands); + Objects.equals(commentCommands, that.commentCommands) && + Objects.equals(maxAnalysisTokenLimit, that.maxAnalysisTokenLimit); } @Override public int hashCode() { return Objects.hash(useLocalMcp, mainBranch, branchAnalysis, ragConfig, - prAnalysisEnabled, branchAnalysisEnabled, installationMethod, commentCommands); + prAnalysisEnabled, branchAnalysisEnabled, installationMethod, + commentCommands, maxAnalysisTokenLimit); } @Override @@ -250,6 +279,7 @@ public String toString() { ", branchAnalysisEnabled=" + branchAnalysisEnabled + ", installationMethod=" + installationMethod + ", commentCommands=" + commentCommands + + ", maxAnalysisTokenLimit=" + maxAnalysisTokenLimit + '}'; } } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/vcs/VcsConnection.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/vcs/VcsConnection.java index 383ea9b1..9c3f231a 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/vcs/VcsConnection.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/model/vcs/VcsConnection.java @@ -116,6 +116,15 @@ public class VcsConnection { @Column(name = "updated_at") private LocalDateTime updatedAt; + /** + * Version field for optimistic locking. + * Prevents concurrent token refresh operations from overwriting each other. + * Initialized to 0 to handle existing records that don't have this field yet. + */ + @Version + @Column(name = "version", nullable = false, columnDefinition = "BIGINT DEFAULT 0") + private Long version = 0L; + /** * Provider-specific configuration (JSON column). * Stores additional settings like OAuth keys for manual connections. diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/analysis/CommentCommandRateLimitRepository.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/analysis/CommentCommandRateLimitRepository.java index 9118391f..35527873 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/analysis/CommentCommandRateLimitRepository.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/analysis/CommentCommandRateLimitRepository.java @@ -31,4 +31,22 @@ int countCommandsInWindow( @Param("projectId") Long projectId, @Param("windowStart") OffsetDateTime windowStart ); + + /** + * Atomic upsert: increments command count if record exists, creates with count=1 if not. + * Uses PostgreSQL ON CONFLICT DO UPDATE to avoid race conditions. + */ + @Modifying + @Query(value = """ + INSERT INTO comment_command_rate_limit (project_id, window_start, command_count, last_command_at) + VALUES (:projectId, :windowStart, 1, NOW()) + ON CONFLICT (project_id, window_start) + DO UPDATE SET + command_count = comment_command_rate_limit.command_count + 1, + last_command_at = NOW() + """, nativeQuery = true) + void upsertCommandCount( + @Param("projectId") Long projectId, + @Param("windowStart") OffsetDateTime windowStart + ); } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/codeanalysis/CodeAnalysisRepository.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/codeanalysis/CodeAnalysisRepository.java index dbab64ca..516ca3c9 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/codeanalysis/CodeAnalysisRepository.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/codeanalysis/CodeAnalysisRepository.java @@ -111,4 +111,52 @@ Page searchAnalyses( }) @Query("SELECT ca FROM CodeAnalysis ca WHERE ca.id = :id") Optional findByIdWithIssues(@Param("id") Long id); + + /** + * Find the most recent ACCEPTED analysis for a project with the same diff fingerprint. + * Used for content-based cache: reuse analysis when the same code changes appear in a different PR. + */ + @org.springframework.data.jpa.repository.EntityGraph(attributePaths = { + "issues", + "project", + "project.workspace", + "project.vcsBinding", + "project.vcsBinding.vcsConnection", + "project.aiBinding" + }) + @Query("SELECT ca FROM CodeAnalysis ca WHERE ca.project.id = :projectId " + + "AND ca.diffFingerprint = :diffFingerprint " + + "AND ca.status = org.rostilos.codecrow.core.model.codeanalysis.AnalysisStatus.ACCEPTED " + + "ORDER BY ca.createdAt DESC LIMIT 1") + Optional findTopByProjectIdAndDiffFingerprint( + @Param("projectId") Long projectId, + @Param("diffFingerprint") String diffFingerprint); + + /** + * Find the most recent ACCEPTED analysis for a project + commit hash (any PR number). + * Fallback cache for close/reopen scenarios where the same commit gets a new PR number. + */ + @org.springframework.data.jpa.repository.EntityGraph(attributePaths = { + "issues", + "project", + "project.workspace", + "project.vcsBinding", + "project.vcsBinding.vcsConnection", + "project.aiBinding" + }) + @Query("SELECT ca FROM CodeAnalysis ca WHERE ca.project.id = :projectId " + + "AND ca.commitHash = :commitHash " + + "AND ca.status = org.rostilos.codecrow.core.model.codeanalysis.AnalysisStatus.ACCEPTED " + + "ORDER BY ca.createdAt DESC LIMIT 1") + Optional findTopByProjectIdAndCommitHash( + @Param("projectId") Long projectId, + @Param("commitHash") String commitHash); + + /** + * Find all analyses for a PR across all versions, ordered by version descending. + * Used to provide LLM with full issue history including resolved issues. + */ + @org.springframework.data.jpa.repository.EntityGraph(attributePaths = {"issues"}) + @Query("SELECT ca FROM CodeAnalysis ca WHERE ca.project.id = :projectId AND ca.prNumber = :prNumber ORDER BY ca.prVersion DESC") + List findAllByProjectIdAndPrNumberOrderByPrVersionDesc(@Param("projectId") Long projectId, @Param("prNumber") Long prNumber); } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/job/JobRepository.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/job/JobRepository.java index 6dd799b5..b7f75a97 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/job/JobRepository.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/persistence/repository/job/JobRepository.java @@ -101,4 +101,8 @@ Page findByProjectIdAndDateRange( @Modifying @Query("DELETE FROM Job j WHERE j.project.id = :projectId") void deleteByProjectId(@Param("projectId") Long projectId); + + @Modifying + @Query("DELETE FROM Job j WHERE j.id = :jobId") + void deleteJobById(@Param("jobId") Long jobId); } diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/CodeAnalysisService.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/CodeAnalysisService.java index 7f3ca41f..f34a58ea 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/CodeAnalysisService.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/CodeAnalysisService.java @@ -24,23 +24,20 @@ @Transactional public class CodeAnalysisService { - private final CodeAnalysisRepository analysisRepository; - private final CodeAnalysisIssueRepository issueRepository; private final CodeAnalysisRepository codeAnalysisRepository; + private final CodeAnalysisIssueRepository issueRepository; private final QualityGateRepository qualityGateRepository; private final QualityGateEvaluator qualityGateEvaluator; private static final Logger log = LoggerFactory.getLogger(CodeAnalysisService.class); @Autowired public CodeAnalysisService( - CodeAnalysisRepository analysisRepository, - CodeAnalysisIssueRepository issueRepository, CodeAnalysisRepository codeAnalysisRepository, + CodeAnalysisIssueRepository issueRepository, QualityGateRepository qualityGateRepository ) { - this.analysisRepository = analysisRepository; - this.issueRepository = issueRepository; this.codeAnalysisRepository = codeAnalysisRepository; + this.issueRepository = issueRepository; this.qualityGateRepository = qualityGateRepository; this.qualityGateEvaluator = new QualityGateEvaluator(); } @@ -54,6 +51,21 @@ public CodeAnalysis createAnalysisFromAiResponse( String commitHash, String vcsAuthorId, String vcsAuthorUsername + ) { + return createAnalysisFromAiResponse(project, analysisData, pullRequestId, + targetBranchName, sourceBranchName, commitHash, vcsAuthorId, vcsAuthorUsername, null); + } + + public CodeAnalysis createAnalysisFromAiResponse( + Project project, + Map analysisData, + Long pullRequestId, + String targetBranchName, + String sourceBranchName, + String commitHash, + String vcsAuthorId, + String vcsAuthorUsername, + String diffFingerprint ) { try { // Check if analysis already exists for this commit (handles webhook retries) @@ -74,6 +86,7 @@ public CodeAnalysis createAnalysisFromAiResponse( analysis.setBranchName(targetBranchName); analysis.setSourceBranchName(sourceBranchName); analysis.setPrVersion(previousVersion + 1); + analysis.setDiffFingerprint(diffFingerprint); return fillAnalysisData(analysis, analysisData, commitHash, vcsAuthorId, vcsAuthorUsername); } catch (Exception e) { @@ -107,9 +120,14 @@ private CodeAnalysis fillAnalysisData( Object issuesObj = analysisData.get("issues"); if (issuesObj == null) { log.warn("No issues found in analysis data"); - return analysisRepository.save(analysis); + return codeAnalysisRepository.save(analysis); } + // Save analysis first to get its ID for resolution tracking + CodeAnalysis savedAnalysis = codeAnalysisRepository.save(analysis); + Long analysisId = savedAnalysis.getId(); + Long prNumber = savedAnalysis.getPrNumber(); + // Handle issues as List (array format from AI) if (issuesObj instanceof List) { List issuesList = (List) issuesObj; @@ -122,9 +140,11 @@ private CodeAnalysis fillAnalysisData( log.warn("Null issue data at index: {}", i); continue; } - CodeAnalysisIssue issue = createIssueFromData(issueData, String.valueOf(i), vcsAuthorId, vcsAuthorUsername); + CodeAnalysisIssue issue = createIssueFromData( + issueData, String.valueOf(i), vcsAuthorId, vcsAuthorUsername, + commitHash, prNumber, analysisId); if (issue != null) { - analysis.addIssue(issue); + savedAnalysis.addIssue(issue); } } catch (Exception e) { log.error("Error processing issue at index '{}': {}", i, e.getMessage(), e); @@ -145,9 +165,11 @@ else if (issuesObj instanceof Map) { continue; } - CodeAnalysisIssue issue = createIssueFromData(issueData, entry.getKey(), vcsAuthorId, vcsAuthorUsername); + CodeAnalysisIssue issue = createIssueFromData( + issueData, entry.getKey(), vcsAuthorId, vcsAuthorUsername, + commitHash, prNumber, analysisId); if (issue != null) { - analysis.addIssue(issue); + savedAnalysis.addIssue(issue); } } catch (Exception e) { log.error("Error processing issue with key '{}': {}", entry.getKey(), e.getMessage(), e); @@ -157,19 +179,24 @@ else if (issuesObj instanceof Map) { log.warn("Issues field is neither List nor Map: {}", issuesObj.getClass().getName()); } - log.info("Successfully created analysis with {} issues", analysis.getIssues().size()); + log.info("Successfully created analysis with {} issues", savedAnalysis.getIssues().size()); - // Evaluate quality gate - QualityGate qualityGate = getQualityGateForAnalysis(analysis); - if (qualityGate != null) { - QualityGateResult qgResult = qualityGateEvaluator.evaluate(analysis, qualityGate); - analysis.setAnalysisResult(qgResult.result()); - log.info("Quality gate '{}' evaluated with result: {}", qualityGate.getName(), qgResult.result()); - } else { - log.info("No quality gate found for analysis, skipping evaluation"); + // Evaluate quality gate — wrapped defensively so a QG failure + // (e.g. detached entity, lazy-init) does not abort the entire analysis + try { + QualityGate qualityGate = getQualityGateForAnalysis(savedAnalysis); + if (qualityGate != null) { + QualityGateResult qgResult = qualityGateEvaluator.evaluate(savedAnalysis, qualityGate); + savedAnalysis.setAnalysisResult(qgResult.result()); + log.info("Quality gate '{}' evaluated with result: {}", qualityGate.getName(), qgResult.result()); + } else { + log.info("No quality gate found for analysis, skipping evaluation"); + } + } catch (Exception qgEx) { + log.warn("Quality gate evaluation failed, analysis will be saved without QG result: {}", qgEx.getMessage()); } - return analysisRepository.save(analysis); + return codeAnalysisRepository.save(savedAnalysis); } catch (Exception e) { log.error("Error creating analysis from AI response: {}", e.getMessage(), e); @@ -221,10 +248,112 @@ public Optional getCodeAnalysisCache(Long projectId, String commit return codeAnalysisRepository.findByProjectIdAndCommitHashAndPrNumber(projectId, commitHash, prNumber).stream().findFirst(); } + /** + * Fallback cache lookup by commit hash only (ignoring PR number). + * Handles close/reopen scenarios where the same commit gets a new PR number. + */ + public Optional getAnalysisByCommitHash(Long projectId, String commitHash) { + return codeAnalysisRepository.findTopByProjectIdAndCommitHash(projectId, commitHash); + } + + /** + * Content-based cache lookup by diff fingerprint. + * Handles branch-cascade flows where the same code changes appear in different PRs + * (e.g. feature→release analyzed, then release→main opens with the same changes). + */ + public Optional getAnalysisByDiffFingerprint(Long projectId, String diffFingerprint) { + if (diffFingerprint == null || diffFingerprint.isBlank()) { + return Optional.empty(); + } + return codeAnalysisRepository.findTopByProjectIdAndDiffFingerprint(projectId, diffFingerprint); + } + + /** + * Clone an existing analysis for a new PR. + * Creates a new CodeAnalysis row with cloned issues, linked to the new PR identity. + * Used when a fingerprint/commit-hash cache hit matches a different PR. + * + * // TODO: Option B — LIGHTWEIGHT mode: instead of full clone, reuse Stage 1 issues + * // but re-run Stage 2 cross-file analysis against the new target branch context. + * // This would catch interaction differences when target branches differ. + * + * // TODO: Consider tracking storage growth from cloned analyses. If it becomes significant, + * // explore referencing the original analysis instead of deep-copying issues. + */ + public CodeAnalysis cloneAnalysisForPr( + CodeAnalysis source, + Project project, + Long newPrNumber, + String commitHash, + String targetBranchName, + String sourceBranchName, + String diffFingerprint + ) { + // Guard against duplicates (same idempotency check as createAnalysisFromAiResponse) + Optional existing = codeAnalysisRepository + .findByProjectIdAndCommitHashAndPrNumber(project.getId(), commitHash, newPrNumber); + if (existing.isPresent()) { + log.info("Cloned analysis already exists for project={}, commit={}, pr={}. Returning existing.", + project.getId(), commitHash, newPrNumber); + return existing.get(); + } + + int previousVersion = codeAnalysisRepository.findMaxPrVersion(project.getId(), newPrNumber).orElse(0); + + CodeAnalysis clone = new CodeAnalysis(); + clone.setProject(project); + clone.setAnalysisType(source.getAnalysisType()); + clone.setPrNumber(newPrNumber); + clone.setCommitHash(commitHash); + clone.setDiffFingerprint(diffFingerprint); + clone.setBranchName(targetBranchName); + clone.setSourceBranchName(sourceBranchName); + clone.setComment(source.getComment()); + clone.setStatus(source.getStatus()); + clone.setAnalysisResult(source.getAnalysisResult()); + clone.setPrVersion(previousVersion + 1); + + // Save first to get an ID + CodeAnalysis saved = codeAnalysisRepository.save(clone); + + // Deep-copy issues + for (CodeAnalysisIssue srcIssue : source.getIssues()) { + CodeAnalysisIssue issueClone = new CodeAnalysisIssue(); + issueClone.setSeverity(srcIssue.getSeverity()); + issueClone.setFilePath(srcIssue.getFilePath()); + issueClone.setLineNumber(srcIssue.getLineNumber()); + issueClone.setReason(srcIssue.getReason()); + issueClone.setSuggestedFixDescription(srcIssue.getSuggestedFixDescription()); + issueClone.setSuggestedFixDiff(srcIssue.getSuggestedFixDiff()); + issueClone.setIssueCategory(srcIssue.getIssueCategory()); + issueClone.setResolved(srcIssue.isResolved()); + issueClone.setResolvedDescription(srcIssue.getResolvedDescription()); + issueClone.setVcsAuthorId(srcIssue.getVcsAuthorId()); + issueClone.setVcsAuthorUsername(srcIssue.getVcsAuthorUsername()); + saved.addIssue(issueClone); + } + + saved.updateIssueCounts(); + CodeAnalysis result = codeAnalysisRepository.save(saved); + log.info("Cloned analysis {} → {} for PR {} (fingerprint={}, {} issues)", + source.getId(), result.getId(), newPrNumber, + diffFingerprint != null ? diffFingerprint.substring(0, 8) + "..." : "null", + result.getIssues().size()); + return result; + } + public Optional getPreviousVersionCodeAnalysis(Long projectId, Long prNumber) { return codeAnalysisRepository.findByProjectIdAndPrNumberWithMaxPrVersion(projectId, prNumber); } + /** + * Get all analyses for a PR across all versions. + * Useful for providing full issue history to AI including resolved issues. + */ + public List getAllPrAnalyses(Long projectId, Long prNumber) { + return codeAnalysisRepository.findAllByProjectIdAndPrNumberOrderByPrVersionDesc(projectId, prNumber); + } + public int getMaxAnalysisPrVersion(Long projectId, Long prNumber) { return codeAnalysisRepository.findMaxPrVersion(projectId, prNumber).orElse(0); } @@ -233,13 +362,43 @@ public Optional findAnalysisByProjectAndPrNumberAndVersion(Long pr return codeAnalysisRepository.findByProjectIdAndPrNumberAndPrVersion(projectId, prNumber, prVersion); } - private CodeAnalysisIssue createIssueFromData(Map issueData, String issueKey, String vcsAuthorId, String vcsAuthorUsername) { + private CodeAnalysisIssue createIssueFromData( + Map issueData, + String issueKey, + String vcsAuthorId, + String vcsAuthorUsername, + String commitHash, + Long prNumber, + Long analysisId + ) { try { CodeAnalysisIssue issue = new CodeAnalysisIssue(); issue.setVcsAuthorId(vcsAuthorId); issue.setVcsAuthorUsername(vcsAuthorUsername); + // Check if this is a persisted issue from previous analysis (has original ID) + Object originalIdObj = issueData.get("id"); + CodeAnalysisIssue originalIssue = null; + if (originalIdObj != null) { + try { + Long originalId = null; + if (originalIdObj instanceof String) { + originalId = Long.parseLong((String) originalIdObj); + } else if (originalIdObj instanceof Number) { + originalId = ((Number) originalIdObj).longValue(); + } + if (originalId != null) { + originalIssue = issueRepository.findById(originalId).orElse(null); + if (originalIssue != null) { + log.debug("Found original issue {} for reconciliation", originalId); + } + } + } catch (NumberFormatException e) { + log.debug("Could not parse issue ID '{}' as Long, treating as new issue", originalIdObj); + } + } + String severityStr = (String) issueData.get("severity"); if (severityStr == null) { log.warn("No severity found for issue {}", issueKey); @@ -299,9 +458,29 @@ private CodeAnalysisIssue createIssueFromData(Map issueData, Str } issue.setSuggestedFixDiff(suggestedFixDiff); - Boolean isResolvedObj = (Boolean) issueData.get("isResolved"); - boolean isResolved = isResolvedObj != null ? isResolvedObj : false; + // Parse isResolved - handle both Boolean and String representations + Object isResolvedObj = issueData.get("isResolved"); + boolean isResolved = false; + if (isResolvedObj instanceof Boolean) { + isResolved = (Boolean) isResolvedObj; + } else if (isResolvedObj instanceof String) { + isResolved = "true".equalsIgnoreCase((String) isResolvedObj); + } issue.setResolved(isResolved); + + log.debug("Issue resolved status: isResolvedObj={}, type={}, parsed={}", + isResolvedObj, isResolvedObj != null ? isResolvedObj.getClass().getSimpleName() : "null", isResolved); + + // If this issue is resolved and we have original issue data, populate resolution tracking + if (isResolved && originalIssue != null) { + issue.setResolvedDescription(reason); // AI provides resolution reason in the 'reason' field + issue.setResolvedByPr(prNumber); + issue.setResolvedCommitHash(commitHash); + issue.setResolvedAnalysisId(analysisId); + issue.setResolvedAt(OffsetDateTime.now()); + issue.setResolvedBy(vcsAuthorUsername); + log.info("Issue {} marked as resolved by PR {} commit {}", originalIdObj, prNumber, commitHash); + } String categoryStr = (String) issueData.get("category"); if (categoryStr != null && !categoryStr.isBlank()) { @@ -310,8 +489,8 @@ private CodeAnalysisIssue createIssueFromData(Map issueData, Str issue.setIssueCategory(IssueCategory.CODE_QUALITY); } - log.debug("Created issue: {} severity, category: {}, file: {}, line: {}", - issue.getSeverity(), issue.getIssueCategory(), issue.getFilePath(), issue.getLineNumber()); + log.debug("Created issue: {} severity, category: {}, file: {}, line: {}, resolved: {}", + issue.getSeverity(), issue.getIssueCategory(), issue.getFilePath(), issue.getLineNumber(), isResolved); return issue; @@ -321,45 +500,52 @@ private CodeAnalysisIssue createIssueFromData(Map issueData, Str } } + /** + * Overload for backward compatibility with callers that don't have resolution context + */ + private CodeAnalysisIssue createIssueFromData(Map issueData, String issueKey, String vcsAuthorId, String vcsAuthorUsername) { + return createIssueFromData(issueData, issueKey, vcsAuthorId, vcsAuthorUsername, null, null, null); + } + public CodeAnalysis createAnalysis(Project project, AnalysisType analysisType) { CodeAnalysis analysis = new CodeAnalysis(); analysis.setProject(project); analysis.setAnalysisType(analysisType); analysis.setStatus(AnalysisStatus.PENDING); - return analysisRepository.save(analysis); + return codeAnalysisRepository.save(analysis); } public CodeAnalysis saveAnalysis(CodeAnalysis analysis) { analysis.updateIssueCounts(); - return analysisRepository.save(analysis); + return codeAnalysisRepository.save(analysis); } public Optional findById(Long id) { - return analysisRepository.findById(id); + return codeAnalysisRepository.findById(id); } public List findByProjectId(Long projectId) { - return analysisRepository.findByProjectIdOrderByCreatedAtDesc(projectId); + return codeAnalysisRepository.findByProjectIdOrderByCreatedAtDesc(projectId); } public List findByProjectIdAndType(Long projectId, AnalysisType analysisType) { - return analysisRepository.findByProjectIdAndAnalysisTypeOrderByCreatedAtDesc(projectId, analysisType); + return codeAnalysisRepository.findByProjectIdAndAnalysisTypeOrderByCreatedAtDesc(projectId, analysisType); } public Optional findByProjectIdAndPrNumber(Long projectId, Long prNumber) { - return analysisRepository.findByProjectIdAndPrNumber(projectId, prNumber); + return codeAnalysisRepository.findByProjectIdAndPrNumber(projectId, prNumber); } public Optional findByProjectIdAndPrNumberAndPrVersion(Long projectId, Long prNumber, int prVersion) { - return analysisRepository.findByProjectIdAndPrNumberAndPrVersion(projectId, prNumber, prVersion); + return codeAnalysisRepository.findByProjectIdAndPrNumberAndPrVersion(projectId, prNumber, prVersion); } public List findByProjectIdAndDateRange(Long projectId, OffsetDateTime startDate, OffsetDateTime endDate) { - return analysisRepository.findByProjectIdAndDateRange(projectId, startDate, endDate); + return codeAnalysisRepository.findByProjectIdAndDateRange(projectId, startDate, endDate); } public List findByProjectIdWithHighSeverityIssues(Long projectId) { - return analysisRepository.findByProjectIdWithHighSeverityIssues(projectId); + return codeAnalysisRepository.findByProjectIdWithHighSeverityIssues(projectId); } /** @@ -377,16 +563,16 @@ public org.springframework.data.domain.Page searchAnalyses( Long prNumber, AnalysisStatus status, org.springframework.data.domain.Pageable pageable) { - return analysisRepository.searchAnalyses(projectId, prNumber, status, pageable); + return codeAnalysisRepository.searchAnalyses(projectId, prNumber, status, pageable); } public Optional findLatestByProjectId(Long projectId) { - return analysisRepository.findLatestByProjectId(projectId); + return codeAnalysisRepository.findLatestByProjectId(projectId); } public AnalysisStats getProjectAnalysisStats(Long projectId) { - long totalAnalyses = analysisRepository.countByProjectId(projectId); - Double avgIssues = analysisRepository.getAverageIssuesPerAnalysis(projectId); + long totalAnalyses = codeAnalysisRepository.countByProjectId(projectId); + Double avgIssues = codeAnalysisRepository.getAverageIssuesPerAnalysis(projectId); long highSeverityCount = issueRepository.countByProjectIdAndSeverity(projectId, IssueSeverity.HIGH); long mediumSeverityCount = issueRepository.countByProjectIdAndSeverity(projectId, IssueSeverity.MEDIUM); @@ -419,11 +605,11 @@ public void markIssueAsResolved(Long issueId) { } public void deleteAnalysis(Long analysisId) { - analysisRepository.deleteById(analysisId); + codeAnalysisRepository.deleteById(analysisId); } public void deleteAllAnalysesByProjectId(Long projectId) { - analysisRepository.deleteByProjectId(projectId); + codeAnalysisRepository.deleteByProjectId(projectId); } public static class AnalysisStats { diff --git a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/JobService.java b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/JobService.java index f239b9c1..3feac5fb 100644 --- a/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/JobService.java +++ b/java-ecosystem/libs/core/src/main/java/org/rostilos/codecrow/core/service/JobService.java @@ -224,6 +224,8 @@ private String getCommandJobTitle(JobType type, Long prNumber) { */ @Transactional public Job startJob(Job job) { + // Re-fetch the job in case it was passed from a different transaction context + job = jobRepository.findById(job.getId()).orElse(job); job.start(); job = jobRepository.save(job); addLog(job, JobLogLevel.INFO, "start", "Job started"); @@ -244,6 +246,8 @@ public Job startJob(String externalId) { */ @Transactional public Job completeJob(Job job) { + // Re-fetch the job in case it was passed from a different transaction context + job = jobRepository.findById(job.getId()).orElse(job); job.complete(); job = jobRepository.save(job); addLog(job, JobLogLevel.INFO, "complete", "Job completed successfully"); @@ -273,10 +277,26 @@ public Job completeJob(Job job, CodeAnalysis codeAnalysis) { */ @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW) public Job failJob(Job job, String errorMessage) { + log.info("failJob called for job {} (id={}), error: {}", + job != null ? job.getExternalId() : "null", + job != null ? job.getId() : "null", + errorMessage); + + if (job == null || job.getId() == null) { + log.error("Cannot fail job - job or job ID is null"); + return job; + } + // Re-fetch the job to ensure we have a fresh entity in this new transaction job = jobRepository.findById(job.getId()).orElse(job); + log.info("Re-fetched job {}, current status: {}", job.getExternalId(), job.getStatus()); + job.fail(errorMessage); + log.info("Job {} status set to FAILED", job.getExternalId()); + job = jobRepository.save(job); + log.info("Job {} saved with status: {}", job.getExternalId(), job.getStatus()); + addLogInNewTransaction(job, JobLogLevel.ERROR, "error", "Job failed: " + errorMessage); notifyJobComplete(job); return job; @@ -287,6 +307,8 @@ public Job failJob(Job job, String errorMessage) { */ @Transactional public Job cancelJob(Job job) { + // Re-fetch the job in case it was passed from a different transaction context + job = jobRepository.findById(job.getId()).orElse(job); job.cancel(); job = jobRepository.save(job); addLog(job, JobLogLevel.WARN, "cancel", "Job cancelled"); @@ -299,6 +321,8 @@ public Job cancelJob(Job job) { */ @Transactional public Job skipJob(Job job, String reason) { + // Re-fetch the job in case it was passed from a different transaction context + job = jobRepository.findById(job.getId()).orElse(job); job.skip(reason); job = jobRepository.save(job); addLog(job, JobLogLevel.INFO, "skipped", reason); @@ -311,14 +335,32 @@ public Job skipJob(Job job, String reason) { * Used for jobs that were created but then determined to be unnecessary * (e.g., branch not matching pattern, PR analysis disabled). * This prevents DB clutter from ignored webhooks. + * Uses REQUIRES_NEW to ensure this runs in its own transaction, + * allowing it to work even if the calling transaction has issues. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW) public void deleteIgnoredJob(Job job, String reason) { log.info("Deleting ignored job {} ({}): {}", job.getExternalId(), job.getJobType(), reason); - // Delete any logs first (foreign key constraint) - jobLogRepository.deleteByJobId(job.getId()); - // Delete the job - jobRepository.delete(job); + Long jobId = job.getId(); + if (jobId == null) { + log.warn("Cannot delete ignored job - job ID is null"); + return; + } + + try { + // Use direct JPQL queries to avoid JPA entity lifecycle issues + // Delete logs first (foreign key constraint) + jobLogRepository.deleteByJobId(jobId); + log.info("Deleted job logs for ignored job {}", job.getExternalId()); + + // Delete the job using direct JPQL query (bypasses entity state tracking) + log.info("About to delete job entity {} (id={})", job.getExternalId(), jobId); + jobRepository.deleteJobById(jobId); + log.info("Successfully deleted ignored job {} from database", job.getExternalId()); + } catch (Exception e) { + log.error("Failed to delete ignored job {}: {} - {}", job.getExternalId(), e.getClass().getSimpleName(), e.getMessage(), e); + throw e; + } } /** @@ -338,8 +380,9 @@ public Job updateProgress(Job job, int progress, String currentStep) { /** * Add a log entry to a job. + * Uses REQUIRES_NEW to commit immediately and not hold the parent transaction open. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW, timeout = 10) public JobLog addLog(Job job, JobLogLevel level, String step, String message) { JobLog logEntry = new JobLog(); logEntry.setJob(job); @@ -383,8 +426,9 @@ public JobLog addLogInNewTransaction(Job job, JobLogLevel level, String step, St /** * Add a log entry with metadata. + * Uses REQUIRES_NEW to commit immediately and not hold the parent transaction open. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW, timeout = 10) public JobLog addLog(Job job, JobLogLevel level, String step, String message, Map metadata) { JobLog logEntry = new JobLog(); logEntry.setJob(job); @@ -407,32 +451,36 @@ public JobLog addLog(Job job, JobLogLevel level, String step, String message, Ma /** * Add an info log. + * Uses REQUIRES_NEW to commit immediately. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW, timeout = 10) public JobLog info(Job job, String step, String message) { return addLog(job, JobLogLevel.INFO, step, message); } /** * Add a warning log. + * Uses REQUIRES_NEW to commit immediately. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW, timeout = 10) public JobLog warn(Job job, String step, String message) { return addLog(job, JobLogLevel.WARN, step, message); } /** * Add an error log. + * Uses REQUIRES_NEW to commit immediately. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW, timeout = 10) public JobLog error(Job job, String step, String message) { return addLog(job, JobLogLevel.ERROR, step, message); } /** * Add a debug log. + * Uses REQUIRES_NEW to commit immediately. */ - @Transactional + @Transactional(propagation = org.springframework.transaction.annotation.Propagation.REQUIRES_NEW, timeout = 10) public JobLog debug(Job job, String step, String message) { return addLog(job, JobLogLevel.DEBUG, step, message); } diff --git a/java-ecosystem/libs/core/src/main/resources/db/migration/1.3.0/V1.3.0__remove_token_limitation_from_ai_connection.sql b/java-ecosystem/libs/core/src/main/resources/db/migration/1.3.0/V1.3.0__remove_token_limitation_from_ai_connection.sql new file mode 100644 index 00000000..fcd0faaf --- /dev/null +++ b/java-ecosystem/libs/core/src/main/resources/db/migration/1.3.0/V1.3.0__remove_token_limitation_from_ai_connection.sql @@ -0,0 +1,5 @@ +-- Remove token_limitation column from ai_connection table +-- Token limitation is now configured per-project in the project configuration JSON +-- Default value is 200000 tokens, configured in ProjectConfig.maxAnalysisTokenLimit + +ALTER TABLE ai_connection DROP COLUMN IF EXISTS token_limitation; diff --git a/java-ecosystem/libs/core/src/main/resources/db/migration/1.3.0/V1.3.0__set_default_vcs_connection_version.sql b/java-ecosystem/libs/core/src/main/resources/db/migration/1.3.0/V1.3.0__set_default_vcs_connection_version.sql new file mode 100644 index 00000000..0fd68ede --- /dev/null +++ b/java-ecosystem/libs/core/src/main/resources/db/migration/1.3.0/V1.3.0__set_default_vcs_connection_version.sql @@ -0,0 +1,5 @@ +-- Add version column if not exists and set default for existing records +ALTER TABLE vcs_connection ADD COLUMN IF NOT EXISTS version BIGINT DEFAULT 0; +UPDATE vcs_connection SET version = 0 WHERE version IS NULL; +ALTER TABLE vcs_connection ALTER COLUMN version SET NOT NULL; +ALTER TABLE vcs_connection ALTER COLUMN version SET DEFAULT 0; \ No newline at end of file diff --git a/java-ecosystem/libs/core/src/main/resources/db/migration/1.4.0/V1.4.0__add_diff_fingerprint_to_code_analysis.sql b/java-ecosystem/libs/core/src/main/resources/db/migration/1.4.0/V1.4.0__add_diff_fingerprint_to_code_analysis.sql new file mode 100644 index 00000000..9a68490a --- /dev/null +++ b/java-ecosystem/libs/core/src/main/resources/db/migration/1.4.0/V1.4.0__add_diff_fingerprint_to_code_analysis.sql @@ -0,0 +1,9 @@ +-- Add diff_fingerprint column for content-based analysis caching. +-- Allows reusing analysis results when the same code changes appear in different PRs +-- (e.g. close/reopen with a new PR number, or branch-cascade flows like feature→release→main). +ALTER TABLE code_analysis ADD COLUMN IF NOT EXISTS diff_fingerprint VARCHAR(64); + +-- Index for fingerprint-based cache lookups: (project_id, diff_fingerprint) +CREATE INDEX IF NOT EXISTS idx_code_analysis_project_diff_fingerprint + ON code_analysis (project_id, diff_fingerprint) + WHERE diff_fingerprint IS NOT NULL; diff --git a/java-ecosystem/libs/core/src/main/resources/db/migration/1.4.0/V1.4.1__deduplicate_branch_issues.sql b/java-ecosystem/libs/core/src/main/resources/db/migration/1.4.0/V1.4.1__deduplicate_branch_issues.sql new file mode 100644 index 00000000..8ce9d6e9 --- /dev/null +++ b/java-ecosystem/libs/core/src/main/resources/db/migration/1.4.0/V1.4.1__deduplicate_branch_issues.sql @@ -0,0 +1,59 @@ +-- V1.4.1: Remove duplicate branch_issue rows that accumulated because the +-- deduplication key was based on code_analysis_issue_id (database PK) rather +-- than on the issue's semantic content. Each PR analysis creates fresh +-- CodeAnalysisIssue rows with new IDs, so the same logical issue +-- (same file, line, severity, category) ended up with N BranchIssue rows. +-- +-- Strategy: for each (branch_id, file_path, line_number, severity, category) +-- group keep only the row with the LOWEST id (oldest / first-detected) and +-- delete the rest. Afterwards recompute the denormalized counters on branch. + +-- 1. Delete duplicate branch_issues, keeping the first (lowest id) per group +DELETE FROM branch_issue +WHERE id NOT IN ( + SELECT keeper_id FROM ( + SELECT MIN(bi.id) AS keeper_id + FROM branch_issue bi + JOIN code_analysis_issue cai ON bi.code_analysis_issue_id = cai.id + GROUP BY bi.branch_id, + cai.file_path, + cai.line_number, + cai.severity, + COALESCE(cai.issue_category, '__NONE__') + ) AS keepers +); + +-- 2. Recompute denormalized branch issue counts +UPDATE branch b SET + total_issues = COALESCE(sub.total_unresolved, 0), + high_severity_count = COALESCE(sub.high_count, 0), + medium_severity_count = COALESCE(sub.medium_count, 0), + low_severity_count = COALESCE(sub.low_count, 0), + info_severity_count = COALESCE(sub.info_count, 0), + resolved_count = COALESCE(sub.resolved_total, 0), + updated_at = NOW() +FROM ( + SELECT + bi.branch_id, + COUNT(*) FILTER (WHERE bi.is_resolved = false) AS total_unresolved, + COUNT(*) FILTER (WHERE bi.is_resolved = false AND cai.severity = 'HIGH') AS high_count, + COUNT(*) FILTER (WHERE bi.is_resolved = false AND cai.severity = 'MEDIUM') AS medium_count, + COUNT(*) FILTER (WHERE bi.is_resolved = false AND cai.severity = 'LOW') AS low_count, + COUNT(*) FILTER (WHERE bi.is_resolved = false AND cai.severity = 'INFO') AS info_count, + COUNT(*) FILTER (WHERE bi.is_resolved = true) AS resolved_total + FROM branch_issue bi + JOIN code_analysis_issue cai ON bi.code_analysis_issue_id = cai.id + GROUP BY bi.branch_id +) sub +WHERE b.id = sub.branch_id; + +-- Also zero-out branches that lost all issues after cleanup +UPDATE branch SET + total_issues = 0, + high_severity_count = 0, + medium_severity_count = 0, + low_severity_count = 0, + info_severity_count = 0, + resolved_count = 0, + updated_at = NOW() +WHERE id NOT IN (SELECT DISTINCT branch_id FROM branch_issue); diff --git a/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTOTest.java b/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTOTest.java index 5a8af6ef..7536258c 100644 --- a/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTOTest.java +++ b/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/ai/AIConnectionDTOTest.java @@ -24,7 +24,7 @@ void shouldCreateWithAllFields() { OffsetDateTime now = OffsetDateTime.now(); AIConnectionDTO dto = new AIConnectionDTO( 1L, "Test Connection", AIProviderKey.ANTHROPIC, "claude-3-opus", - now, now, 100000 + now, now ); assertThat(dto.id()).isEqualTo(1L); @@ -33,14 +33,13 @@ void shouldCreateWithAllFields() { assertThat(dto.aiModel()).isEqualTo("claude-3-opus"); assertThat(dto.createdAt()).isEqualTo(now); assertThat(dto.updatedAt()).isEqualTo(now); - assertThat(dto.tokenLimitation()).isEqualTo(100000); } @Test @DisplayName("should create AIConnectionDTO with null optional fields") void shouldCreateWithNullOptionalFields() { AIConnectionDTO dto = new AIConnectionDTO( - 1L, null, AIProviderKey.OPENAI, null, null, null, 50000 + 1L, null, AIProviderKey.OPENAI, null, null, null ); assertThat(dto.id()).isEqualTo(1L); @@ -53,24 +52,14 @@ void shouldCreateWithNullOptionalFields() { @Test @DisplayName("should create AIConnectionDTO with different providers") void shouldCreateWithDifferentProviders() { - AIConnectionDTO openai = new AIConnectionDTO(1L, "OpenAI", AIProviderKey.OPENAI, "gpt-4", null, null, 100000); - AIConnectionDTO anthropic = new AIConnectionDTO(2L, "Anthropic", AIProviderKey.ANTHROPIC, "claude-3", null, null, 200000); - AIConnectionDTO google = new AIConnectionDTO(3L, "Google", AIProviderKey.GOOGLE, "gemini-pro", null, null, 150000); + AIConnectionDTO openai = new AIConnectionDTO(1L, "OpenAI", AIProviderKey.OPENAI, "gpt-4", null, null); + AIConnectionDTO anthropic = new AIConnectionDTO(2L, "Anthropic", AIProviderKey.ANTHROPIC, "claude-3", null, null); + AIConnectionDTO google = new AIConnectionDTO(3L, "Google", AIProviderKey.GOOGLE, "gemini-pro", null, null); assertThat(openai.providerKey()).isEqualTo(AIProviderKey.OPENAI); assertThat(anthropic.providerKey()).isEqualTo(AIProviderKey.ANTHROPIC); assertThat(google.providerKey()).isEqualTo(AIProviderKey.GOOGLE); } - - @Test - @DisplayName("should support different token limitations") - void shouldSupportDifferentTokenLimitations() { - AIConnectionDTO small = new AIConnectionDTO(1L, "Small", AIProviderKey.OPENAI, "gpt-3.5", null, null, 16000); - AIConnectionDTO large = new AIConnectionDTO(2L, "Large", AIProviderKey.ANTHROPIC, "claude-3", null, null, 200000); - - assertThat(small.tokenLimitation()).isEqualTo(16000); - assertThat(large.tokenLimitation()).isEqualTo(200000); - } } @Nested @@ -85,7 +74,6 @@ void shouldConvertWithAllFields() { connection.setName("Production AI"); setField(connection, "providerKey", AIProviderKey.ANTHROPIC); setField(connection, "aiModel", "claude-3-opus"); - setField(connection, "tokenLimitation", 100000); AIConnectionDTO dto = AIConnectionDTO.fromAiConnection(connection); @@ -93,7 +81,6 @@ void shouldConvertWithAllFields() { assertThat(dto.name()).isEqualTo("Production AI"); assertThat(dto.providerKey()).isEqualTo(AIProviderKey.ANTHROPIC); assertThat(dto.aiModel()).isEqualTo("claude-3-opus"); - assertThat(dto.tokenLimitation()).isEqualTo(100000); } @Test @@ -104,7 +91,6 @@ void shouldConvertWithNullName() { connection.setName(null); setField(connection, "providerKey", AIProviderKey.OPENAI); setField(connection, "aiModel", "gpt-4"); - setField(connection, "tokenLimitation", 50000); AIConnectionDTO dto = AIConnectionDTO.fromAiConnection(connection); @@ -120,7 +106,6 @@ void shouldConvertWithNullModel() { connection.setName("Test"); setField(connection, "providerKey", AIProviderKey.GOOGLE); setField(connection, "aiModel", null); - setField(connection, "tokenLimitation", 75000); AIConnectionDTO dto = AIConnectionDTO.fromAiConnection(connection); @@ -134,7 +119,6 @@ void shouldConvertAllProviderTypes() { AIConnection connection = new AIConnection(); setField(connection, "id", 1L); setField(connection, "providerKey", providerKey); - setField(connection, "tokenLimitation", 100000); AIConnectionDTO dto = AIConnectionDTO.fromAiConnection(connection); @@ -149,7 +133,6 @@ void shouldHandleTimestamps() { setField(connection, "id", 1L); connection.setName("Test"); setField(connection, "providerKey", AIProviderKey.ANTHROPIC); - setField(connection, "tokenLimitation", 100000); AIConnectionDTO dto = AIConnectionDTO.fromAiConnection(connection); @@ -165,8 +148,8 @@ class EqualityTests { @DisplayName("should be equal for same values") void shouldBeEqualForSameValues() { OffsetDateTime now = OffsetDateTime.now(); - AIConnectionDTO dto1 = new AIConnectionDTO(1L, "Test", AIProviderKey.OPENAI, "gpt-4", now, now, 100000); - AIConnectionDTO dto2 = new AIConnectionDTO(1L, "Test", AIProviderKey.OPENAI, "gpt-4", now, now, 100000); + AIConnectionDTO dto1 = new AIConnectionDTO(1L, "Test", AIProviderKey.OPENAI, "gpt-4", now, now); + AIConnectionDTO dto2 = new AIConnectionDTO(1L, "Test", AIProviderKey.OPENAI, "gpt-4", now, now); assertThat(dto1).isEqualTo(dto2); assertThat(dto1.hashCode()).isEqualTo(dto2.hashCode()); @@ -176,8 +159,8 @@ void shouldBeEqualForSameValues() { @DisplayName("should not be equal for different values") void shouldNotBeEqualForDifferentValues() { OffsetDateTime now = OffsetDateTime.now(); - AIConnectionDTO dto1 = new AIConnectionDTO(1L, "Test1", AIProviderKey.OPENAI, "gpt-4", now, now, 100000); - AIConnectionDTO dto2 = new AIConnectionDTO(2L, "Test2", AIProviderKey.ANTHROPIC, "claude", now, now, 200000); + AIConnectionDTO dto1 = new AIConnectionDTO(1L, "Test1", AIProviderKey.OPENAI, "gpt-4", now, now); + AIConnectionDTO dto2 = new AIConnectionDTO(2L, "Test2", AIProviderKey.ANTHROPIC, "claude", now, now); assertThat(dto1).isNotEqualTo(dto2); } diff --git a/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/project/ProjectDTOTest.java b/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/project/ProjectDTOTest.java index 63096a11..2fe071a6 100644 --- a/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/project/ProjectDTOTest.java +++ b/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/dto/project/ProjectDTOTest.java @@ -50,7 +50,7 @@ void shouldCreateWithAllFields() { 20L, "namespace", "main", "main", 100L, stats, ragConfig, true, false, "WEBHOOK", - commandsConfig, true, 50L + commandsConfig, true, 50L, 200000 ); assertThat(dto.id()).isEqualTo(1L); @@ -84,7 +84,7 @@ void shouldCreateWithNullOptionalFields() { 1L, "Test", null, true, null, null, null, null, null, null, null, null, null, null, null, - null, null, null, null, null, null, null + null, null, null, null, null, null, null, null ); assertThat(dto.description()).isNull(); diff --git a/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/model/ai/AIConnectionTest.java b/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/model/ai/AIConnectionTest.java index d2dd3d4d..6adcfaa6 100644 --- a/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/model/ai/AIConnectionTest.java +++ b/java-ecosystem/libs/core/src/test/java/org/rostilos/codecrow/core/model/ai/AIConnectionTest.java @@ -67,25 +67,12 @@ void shouldSetAndGetApiKeyEncrypted() { aiConnection.setApiKeyEncrypted("encrypted-api-key-xyz"); assertThat(aiConnection.getApiKeyEncrypted()).isEqualTo("encrypted-api-key-xyz"); } - - @Test - @DisplayName("Should set and get tokenLimitation") - void shouldSetAndGetTokenLimitation() { - aiConnection.setTokenLimitation(50000); - assertThat(aiConnection.getTokenLimitation()).isEqualTo(50000); - } } @Nested @DisplayName("Default value tests") class DefaultValueTests { - @Test - @DisplayName("Default tokenLimitation should be 100000") - void defaultTokenLimitationShouldBe100000() { - assertThat(aiConnection.getTokenLimitation()).isEqualTo(100000); - } - @Test @DisplayName("Id should be null for new entity") void idShouldBeNullForNewEntity() { @@ -154,14 +141,12 @@ void shouldBeAbleToUpdateAllFields() { aiConnection.setProviderKey(AIProviderKey.ANTHROPIC); aiConnection.setAiModel("claude-3-opus"); aiConnection.setApiKeyEncrypted("new-encrypted-key"); - aiConnection.setTokenLimitation(200000); assertThat(aiConnection.getName()).isEqualTo("Updated Name"); assertThat(aiConnection.getWorkspace()).isSameAs(workspace); assertThat(aiConnection.getProviderKey()).isEqualTo(AIProviderKey.ANTHROPIC); assertThat(aiConnection.getAiModel()).isEqualTo("claude-3-opus"); assertThat(aiConnection.getApiKeyEncrypted()).isEqualTo("new-encrypted-key"); - assertThat(aiConnection.getTokenLimitation()).isEqualTo(200000); } @Test diff --git a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/client/RagPipelineClient.java b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/client/RagPipelineClient.java index efb34d84..44456325 100644 --- a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/client/RagPipelineClient.java +++ b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/client/RagPipelineClient.java @@ -22,16 +22,19 @@ public class RagPipelineClient { private final ObjectMapper objectMapper; private final String ragApiUrl; private final boolean ragEnabled; + private final String serviceSecret; public RagPipelineClient( @Value("${codecrow.rag.api.url:http://rag-pipeline:8001}") String ragApiUrl, @Value("${codecrow.rag.api.enabled:false}") boolean ragEnabled, @Value("${codecrow.rag.api.timeout.connect:30}") int connectTimeout, @Value("${codecrow.rag.api.timeout.read:120}") int readTimeout, - @Value("${codecrow.rag.api.timeout.indexing:14400}") int indexingTimeout + @Value("${codecrow.rag.api.timeout.indexing:14400}") int indexingTimeout, + @Value("${codecrow.rag.api.secret:}") String serviceSecret ) { this.ragApiUrl = ragApiUrl; this.ragEnabled = ragEnabled; + this.serviceSecret = serviceSecret != null ? serviceSecret : ""; this.httpClient = new OkHttpClient.Builder() .connectTimeout(connectTimeout, java.util.concurrent.TimeUnit.SECONDS) @@ -209,10 +212,11 @@ public void deleteIndex(String workspace, String project, String branch) throws } String url = String.format("%s/index/%s/%s/%s", ragApiUrl, workspace, project, branch); - Request request = new Request.Builder() + Request.Builder builder = new Request.Builder() .url(url) - .delete() - .build(); + .delete(); + addAuthHeader(builder); + Request request = builder.build(); try (Response response = httpClient.newCall(request).execute()) { if (!response.isSuccessful()) { @@ -241,10 +245,11 @@ public boolean deleteBranch(String workspace, String project, String branch) thr String encodedBranch = java.net.URLEncoder.encode(branch, java.nio.charset.StandardCharsets.UTF_8); String url = String.format("%s/index/%s/%s/branch/%s", ragApiUrl, workspace, project, encodedBranch); - Request request = new Request.Builder() + Request.Builder builder = new Request.Builder() .url(url) - .delete() - .build(); + .delete(); + addAuthHeader(builder); + Request request = builder.build(); try (Response response = httpClient.newCall(request).execute()) { if (response.isSuccessful()) { @@ -271,10 +276,11 @@ public List getIndexedBranches(String workspace, String project) { try { String url = String.format("%s/index/%s/%s/branches", ragApiUrl, workspace, project); - Request request = new Request.Builder() + Request.Builder builder = new Request.Builder() .url(url) - .get() - .build(); + .get(); + addAuthHeader(builder); + Request request = builder.build(); try (Response response = httpClient.newCall(request).execute()) { if (response.isSuccessful() && response.body() != null) { @@ -311,10 +317,11 @@ public List> getIndexedBranchesWithStats(String workspace, S try { String url = String.format("%s/index/%s/%s/branches", ragApiUrl, workspace, project); - Request request = new Request.Builder() + Request.Builder builder = new Request.Builder() .url(url) - .get() - .build(); + .get(); + addAuthHeader(builder); + Request request = builder.build(); try (Response response = httpClient.newCall(request).execute()) { if (response.isSuccessful() && response.body() != null) { @@ -376,10 +383,11 @@ public boolean isHealthy() { } try { - Request request = new Request.Builder() + Request.Builder builder = new Request.Builder() .url(ragApiUrl + "/health") - .get() - .build(); + .get(); + addAuthHeader(builder); + Request request = builder.build(); try (Response response = httpClient.newCall(request).execute()) { return response.isSuccessful(); @@ -398,15 +406,25 @@ private Map postLongRunning(String url, Map payl return doRequest(url, payload, longRunningHttpClient); } + /** + * Adds the x-service-secret header to the request if a secret is configured. + */ + private void addAuthHeader(Request.Builder builder) { + if (!serviceSecret.isEmpty()) { + builder.addHeader("x-service-secret", serviceSecret); + } + } + @SuppressWarnings("unchecked") private Map doRequest(String url, Map payload, OkHttpClient client) throws IOException { String json = objectMapper.writeValueAsString(payload); RequestBody body = RequestBody.create(json, JSON); - Request request = new Request.Builder() + Request.Builder builder = new Request.Builder() .url(url) - .post(body) - .build(); + .post(body); + addAuthHeader(builder); + Request request = builder.build(); try (Response response = client.newCall(request).execute()) { String responseBody = response.body() != null ? response.body().string() : "{}"; diff --git a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/IncrementalRagUpdateService.java b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/IncrementalRagUpdateService.java index bae4b707..04a6fd5f 100644 --- a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/IncrementalRagUpdateService.java +++ b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/IncrementalRagUpdateService.java @@ -43,15 +43,29 @@ public IncrementalRagUpdateService( public boolean shouldPerformIncrementalUpdate(Project project) { if (!ragApiEnabled) { + log.info("shouldPerformIncrementalUpdate: ragApiEnabled=false for project={}", project.getId()); return false; } ProjectConfig config = project.getConfiguration(); - if (config == null || config.ragConfig() == null || !config.ragConfig().enabled()) { + if (config == null) { + log.info("shouldPerformIncrementalUpdate: config is null for project={}", project.getId()); + return false; + } + + if (config.ragConfig() == null) { + log.info("shouldPerformIncrementalUpdate: ragConfig is null for project={}", project.getId()); + return false; + } + + if (!config.ragConfig().enabled()) { + log.info("shouldPerformIncrementalUpdate: ragConfig.enabled=false for project={}", project.getId()); return false; } - return ragIndexTrackingService.isProjectIndexed(project); + boolean isIndexed = ragIndexTrackingService.isProjectIndexed(project); + log.info("shouldPerformIncrementalUpdate: project={} isProjectIndexed={}", project.getId(), isIndexed); + return isIndexed; } public Map performIncrementalUpdate( diff --git a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/RagOperationsServiceImpl.java b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/RagOperationsServiceImpl.java index 65720a4d..56c54d09 100644 --- a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/RagOperationsServiceImpl.java +++ b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/RagOperationsServiceImpl.java @@ -387,6 +387,8 @@ public boolean ensureBranchIndexForPrTarget( // With single-collection architecture, we check if branch has any indexed data // If not, we need to index the branch + log.info("ensureBranchIndexForPrTarget called for project={}, branch={}", project.getId(), targetBranch); + if (!isRagEnabled(project)) { log.debug("RAG not enabled for project={}", project.getId()); return false; @@ -402,12 +404,6 @@ public boolean ensureBranchIndexForPrTarget( return false; } - // Check if branch already has indexed data - if (isBranchIndexReady(project, targetBranch)) { - log.debug("Branch {} already indexed for project={}", targetBranch, project.getId()); - return true; - } - // Get VCS connection info VcsRepoBinding vcsRepoBinding = project.getVcsRepoBinding(); if (vcsRepoBinding == null) { @@ -424,43 +420,57 @@ public boolean ensureBranchIndexForPrTarget( // Same branch? Already indexed via main index if (targetBranch.equals(baseBranch)) { - log.debug("Target branch is same as base branch - already indexed"); + log.debug("Target branch {} is same as base branch {} - already indexed", targetBranch, baseBranch); return true; } + // Check if branch already has indexed data (RagBranchIndex exists) + // Note: We still proceed with diff check to ensure any new changes are indexed + boolean branchIndexExists = isBranchIndexReady(project, targetBranch); + log.info("Branch index status for project={}, branch={}: exists={}", + project.getId(), targetBranch, branchIndexExists); + try { - log.info("Indexing branch data for project={}, branch={}", project.getId(), targetBranch); + log.info("Fetching diff between base branch '{}' and target branch '{}' for project={}", + baseBranch, targetBranch, project.getId()); eventConsumer.accept(Map.of( "type", "status", "state", "branch_index", - "message", String.format("Indexing branch '%s'", targetBranch) + "message", String.format("Indexing branch '%s' (diff vs '%s')", targetBranch, baseBranch) )); // Fetch diff between base branch and target branch VcsClient vcsClient = vcsClientProvider.getClient(vcsConnection); String rawDiff = vcsClient.getBranchDiff(workspaceSlug, repoSlug, baseBranch, targetBranch); + log.info("Branch diff result for project={}, branch={}: diffLength={}", + project.getId(), targetBranch, rawDiff != null ? rawDiff.length() : 0); + if (rawDiff == null || rawDiff.isEmpty()) { - log.debug("No diff between {} and {} - using main index", baseBranch, targetBranch); + log.info("No diff between '{}' and '{}' - branch has same content as base, using main index", + baseBranch, targetBranch); eventConsumer.accept(Map.of( "type", "info", - "message", String.format("No changes between %s and %s", baseBranch, targetBranch) + "message", String.format("No changes between %s and %s - using main branch index", baseBranch, targetBranch) )); return true; } // Get latest commit hash on target branch String targetCommit = vcsClient.getLatestCommitHash(workspaceSlug, repoSlug, targetBranch); + log.info("Target branch '{}' commit hash: {}", targetBranch, targetCommit); // Trigger incremental update for this branch + log.info("Triggering incremental update for project={}, branch={}, commit={}, diffBytes={}", + project.getId(), targetBranch, targetCommit, rawDiff.length()); triggerIncrementalUpdate(project, targetBranch, targetCommit, rawDiff, eventConsumer); return true; } catch (Exception e) { - log.error("Failed to index branch data for project={}, branch={}", - project.getId(), targetBranch, e); + log.error("Failed to index branch data for project={}, branch={}: {}", + project.getId(), targetBranch, e.getMessage(), e); eventConsumer.accept(Map.of( "type", "warning", "state", "branch_error", @@ -638,8 +648,10 @@ public boolean ensureRagIndexUpToDate( String targetBranch, Consumer> eventConsumer ) { + log.info("ensureRagIndexUpToDate called for project={}, targetBranch={}", project.getId(), targetBranch); + if (!isRagEnabled(project)) { - log.debug("RAG not enabled for project={}", project.getId()); + log.info("RAG not enabled for project={}", project.getId()); return false; } @@ -656,16 +668,20 @@ public boolean ensureRagIndexUpToDate( // Get base branch (main branch) String baseBranch = getBaseBranch(project); + log.info("Base branch for project={}: '{}'", project.getId(), baseBranch); try { VcsClient vcsClient = vcsClientProvider.getClient(vcsConnection); // Case 1: Target branch is the main branch - check/update main RAG index if (targetBranch.equals(baseBranch)) { + log.info("Target branch '{}' equals base branch '{}' - updating main index only", targetBranch, baseBranch); return ensureMainIndexUpToDate(project, targetBranch, vcsClient, workspaceSlug, repoSlug, eventConsumer); } // Case 2: Different branch - ensure main index is ready, then ensure branch is indexed + log.info("Target branch '{}' differs from base branch '{}' - will ensure branch index", targetBranch, baseBranch); + // First ensure main index is up to date ensureMainIndexUpToDate(project, baseBranch, vcsClient, workspaceSlug, repoSlug, eventConsumer); @@ -745,6 +761,7 @@ private boolean ensureMainIndexUpToDate( /** * Ensures the branch index is up-to-date with the current commit. + * For non-main branches, this compares against the previously indexed commit. */ private boolean ensureBranchIndexUpToDate( Project project, @@ -755,38 +772,49 @@ private boolean ensureBranchIndexUpToDate( String repoSlug, Consumer> eventConsumer ) throws IOException { + log.info("ensureBranchIndexUpToDate called for project={}, targetBranch={}, baseBranch={}", + project.getId(), targetBranch, baseBranch); + // Get current commit on target branch String currentCommit = vcsClient.getLatestCommitHash(workspaceSlug, repoSlug, targetBranch); + log.info("Current commit on branch '{}': {}", targetBranch, currentCommit); // Check if we have branch index tracking Optional branchIndexOpt = ragBranchIndexRepository .findByProjectIdAndBranchName(project.getId(), targetBranch); if (branchIndexOpt.isEmpty()) { - // No branch index exists - create it - log.info("Branch index does not exist for project={}, branch={} - creating", - project.getId(), targetBranch); + // No branch index exists - create it by getting full diff vs main + log.info("No RagBranchIndex entry found for project={}, branch={} - will create with full diff vs {}", + project.getId(), targetBranch, baseBranch); return ensureBranchIndexForPrTarget(project, targetBranch, eventConsumer); } RagBranchIndex branchIndex = branchIndexOpt.get(); String indexedCommit = branchIndex.getCommitHash(); + log.info("Existing RagBranchIndex for project={}, branch={}: indexedCommit={}", + project.getId(), targetBranch, indexedCommit); // If commits match, index is up to date if (currentCommit.equals(indexedCommit)) { - log.debug("Branch index is up-to-date for project={}, branch={}, commit={}", + log.info("Branch index is up-to-date for project={}, branch={}, commit={}", project.getId(), targetBranch, currentCommit); return true; } - log.info("Branch index outdated for project={}, branch={}: indexed={}, current={}", + // Branch index exists but commit changed - do INCREMENTAL update (only new changes) + // Full diff vs main is only done on INITIAL indexing (when RagBranchIndex doesn't exist) + log.info("Branch index outdated for project={}, branch={}: indexed={}, current={} - fetching incremental diff", project.getId(), targetBranch, indexedCommit, currentCommit); - // Fetch diff between indexed commit and current commit on this branch + // Fetch diff between last indexed commit and current commit (incremental) String rawDiff = vcsClient.getBranchDiff(workspaceSlug, repoSlug, indexedCommit, currentCommit); + log.info("Incremental diff for branch '{}' ({} -> {}): bytes={}", + targetBranch, indexedCommit.substring(0, 7), currentCommit.substring(0, 7), + rawDiff != null ? rawDiff.length() : 0); if (rawDiff == null || rawDiff.isEmpty()) { - log.debug("No diff between {} and {} - index is up to date", indexedCommit, currentCommit); + log.info("No diff between {} and {} - updating commit hash only", indexedCommit, currentCommit); // Update commit hash branchIndex.setCommitHash(currentCommit); branchIndex.setUpdatedAt(OffsetDateTime.now()); @@ -797,12 +825,12 @@ private boolean ensureBranchIndexUpToDate( eventConsumer.accept(Map.of( "type", "status", "state", "branch_update", - "message", String.format("Updating branch %s index from %s to %s", - targetBranch, indexedCommit.substring(0, Math.min(7, indexedCommit.length())), - currentCommit.substring(0, 7)) + "message", String.format("Updating branch %s index (incremental: %d bytes)", targetBranch, rawDiff.length()) )); // Trigger incremental update for this branch + log.info("Triggering incremental branch update for '{}' with {} bytes diff", + targetBranch, rawDiff.length()); triggerIncrementalUpdate(project, targetBranch, currentCommit, rawDiff, eventConsumer); return true; diff --git a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/VcsRagIndexingService.java b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/VcsRagIndexingService.java index 6c05e098..e8edbdf9 100644 --- a/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/VcsRagIndexingService.java +++ b/java-ecosystem/libs/rag-engine/src/main/java/org/rostilos/codecrow/ragengine/service/VcsRagIndexingService.java @@ -98,40 +98,37 @@ public Map indexProjectFromVcs( String branch = determineBranch(requestBranch, config); - Job job = jobService.createRagIndexJob(project, null); - if (job != null) { - jobService.startJob(job); - jobService.logToJob(job, JobLogLevel.INFO, "init", - "Starting RAG indexing for branch: " + branch); - } - - messageConsumer.accept(Map.of( - "type", "progress", - "stage", "init", - "message", "Starting RAG indexing for branch: " + branch - )); - + // Check if indexing can start BEFORE creating job to avoid orphan "failed" jobs if (!ragIndexTrackingService.canStartIndexing(project)) { log.warn("RAG indexing already in progress for project: {}", project.getName()); - if (job != null) { - jobService.failJob(job, "RAG indexing is already in progress"); - } return Map.of("status", "locked", "message", "RAG indexing is already in progress"); } + // Try to acquire lock BEFORE creating job - this is the authoritative check Optional lockKey = analysisLockService.acquireLock( project, branch, AnalysisLockType.RAG_INDEXING ); if (lockKey.isEmpty()) { - log.warn("Failed to acquire RAG indexing lock for project: {}", project.getName()); - if (job != null) { - jobService.failJob(job, "Could not acquire lock for RAG indexing"); - } - return Map.of("status", "locked", "message", "Could not acquire lock for RAG indexing"); + log.warn("Failed to acquire RAG indexing lock for project: {} (another process holds the lock)", project.getName()); + return Map.of("status", "locked", "message", "RAG indexing is already in progress (lock held by another process)"); } try { + // Now that we have the lock, create and start the job + Job job = jobService.createRagIndexJob(project, null); + if (job != null) { + jobService.startJob(job); + jobService.logToJob(job, JobLogLevel.INFO, "init", + "Starting RAG indexing for branch: " + branch); + } + + messageConsumer.accept(Map.of( + "type", "progress", + "stage", "init", + "message", "Starting RAG indexing for branch: " + branch + )); + return performIndexing(project, vcsConnection, workspaceSlug, repoSlug, branch, config, messageConsumer, job); } finally { analysisLockService.releaseLock(lockKey.get()); diff --git a/java-ecosystem/libs/rag-engine/src/test/java/org/rostilos/codecrow/ragengine/client/RagPipelineClientTest.java b/java-ecosystem/libs/rag-engine/src/test/java/org/rostilos/codecrow/ragengine/client/RagPipelineClientTest.java index 274fc1c5..f2694048 100644 --- a/java-ecosystem/libs/rag-engine/src/test/java/org/rostilos/codecrow/ragengine/client/RagPipelineClientTest.java +++ b/java-ecosystem/libs/rag-engine/src/test/java/org/rostilos/codecrow/ragengine/client/RagPipelineClientTest.java @@ -33,7 +33,8 @@ void setUp() throws IOException { true, // enabled 5, // connect timeout 10, // read timeout - 20 // indexing timeout + 20, // indexing timeout + "test-secret" // service secret ); objectMapper = new ObjectMapper(); @@ -71,7 +72,7 @@ void testDeleteFiles_WhenDisabled() throws Exception { RagPipelineClient disabledClient = new RagPipelineClient( mockWebServer.url("/").toString(), false, // disabled - 5, 10, 20 + 5, 10, 20, "" ); List files = List.of("file1.java"); @@ -123,7 +124,7 @@ void testSemanticSearch_WhenDisabled() throws Exception { RagPipelineClient disabledClient = new RagPipelineClient( mockWebServer.url("/").toString(), false, - 5, 10, 20 + 5, 10, 20, "" ); Map result = disabledClient.semanticSearch( @@ -159,7 +160,7 @@ void testGetPRContext_WhenDisabled() throws Exception { RagPipelineClient disabledClient = new RagPipelineClient( mockWebServer.url("/").toString(), false, - 5, 10, 20 + 5, 10, 20, "" ); Map result = disabledClient.getPRContext( @@ -226,7 +227,7 @@ void testUpdateFiles_WhenDisabled() throws Exception { RagPipelineClient disabledClient = new RagPipelineClient( mockWebServer.url("/").toString(), false, - 5, 10, 20 + 5, 10, 20, "" ); Map result = disabledClient.updateFiles( @@ -243,7 +244,8 @@ void testConstructor_WithDefaults() { true, 30, 120, - 14400 + 14400, + "" ); assertThat(defaultClient).isNotNull(); @@ -268,4 +270,35 @@ void testNetworkError_ThrowsIOException() throws IOException { )) .isInstanceOf(IOException.class); } + + @Test + void testServiceSecretHeader_SentOnRequests() throws Exception { + Map mockResponse = Map.of("status", "success"); + mockWebServer.enqueue(new MockResponse() + .setBody(objectMapper.writeValueAsString(mockResponse)) + .addHeader("Content-Type", "application/json")); + + client.deleteFiles(List.of("file.java"), "ws", "proj", "main"); + + RecordedRequest request = mockWebServer.takeRequest(); + assertThat(request.getHeader("x-service-secret")).isEqualTo("test-secret"); + } + + @Test + void testServiceSecretHeader_NotSentWhenEmpty() throws Exception { + RagPipelineClient noSecretClient = new RagPipelineClient( + mockWebServer.url("/").toString(), + true, 5, 10, 20, "" + ); + + Map mockResponse = Map.of("status", "success"); + mockWebServer.enqueue(new MockResponse() + .setBody(objectMapper.writeValueAsString(mockResponse)) + .addHeader("Content-Type", "application/json")); + + noSecretClient.deleteFiles(List.of("file.java"), "ws", "proj", "main"); + + RecordedRequest request = mockWebServer.takeRequest(); + assertThat(request.getHeader("x-service-secret")).isNull(); + } } diff --git a/java-ecosystem/libs/security/src/main/java/module-info.java b/java-ecosystem/libs/security/src/main/java/module-info.java index 846a95bc..b23eadb4 100644 --- a/java-ecosystem/libs/security/src/main/java/module-info.java +++ b/java-ecosystem/libs/security/src/main/java/module-info.java @@ -19,4 +19,5 @@ exports org.rostilos.codecrow.security.jwt.utils; exports org.rostilos.codecrow.security.pipelineagent; exports org.rostilos.codecrow.security.pipelineagent.jwt; + exports org.rostilos.codecrow.security.annotations; } diff --git a/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/HasOwnerOrAdminRights.java b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/HasOwnerOrAdminRights.java new file mode 100644 index 00000000..02b28f62 --- /dev/null +++ b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/HasOwnerOrAdminRights.java @@ -0,0 +1,13 @@ +package org.rostilos.codecrow.security.annotations; + +import org.springframework.security.access.prepost.PreAuthorize; +import java.lang.annotation.*; + +@Target({ElementType.METHOD, ElementType.TYPE}) +@Retention(RetentionPolicy.RUNTIME) +@Inherited +@Documented +@PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") +public @interface HasOwnerOrAdminRights { +} + diff --git a/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceMember.java b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceMember.java new file mode 100644 index 00000000..0a5c3b99 --- /dev/null +++ b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceMember.java @@ -0,0 +1,12 @@ +package org.rostilos.codecrow.security.annotations; + +import org.springframework.security.access.prepost.PreAuthorize; +import java.lang.annotation.*; + +@Target({ElementType.METHOD, ElementType.TYPE}) +@Retention(RetentionPolicy.RUNTIME) +@Inherited +@Documented +@PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") +public @interface IsWorkspaceMember { +} \ No newline at end of file diff --git a/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceOwner.java b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceOwner.java new file mode 100644 index 00000000..40af56ad --- /dev/null +++ b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceOwner.java @@ -0,0 +1,12 @@ +package org.rostilos.codecrow.security.annotations; + +import org.springframework.security.access.prepost.PreAuthorize; +import java.lang.annotation.*; + +@Target({ElementType.METHOD, ElementType.TYPE}) +@Retention(RetentionPolicy.RUNTIME) +@Inherited +@Documented +@PreAuthorize("@workspaceSecurity.isWorkspaceOwner(#workspaceSlug, authentication)") +public @interface IsWorkspaceOwner { +} diff --git a/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceReviewer.java b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceReviewer.java new file mode 100644 index 00000000..c6925de2 --- /dev/null +++ b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/annotations/IsWorkspaceReviewer.java @@ -0,0 +1,12 @@ +package org.rostilos.codecrow.security.annotations; + +import org.springframework.security.access.prepost.PreAuthorize; +import java.lang.annotation.*; + +@Target({ElementType.METHOD, ElementType.TYPE}) +@Retention(RetentionPolicy.RUNTIME) +@Inherited +@Documented +@PreAuthorize("@workspaceSecurity.isWorkspaceReviewer(#workspaceSlug, authentication)") +public @interface IsWorkspaceReviewer { +} diff --git a/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/web/WebSecurityConfig.java b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/web/WebSecurityConfig.java index 77a284f3..10a65065 100644 --- a/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/web/WebSecurityConfig.java +++ b/java-ecosystem/libs/security/src/main/java/org/rostilos/codecrow/security/web/WebSecurityConfig.java @@ -1,5 +1,6 @@ package org.rostilos.codecrow.security.web; +import jakarta.servlet.DispatcherType; import java.util.Arrays; import org.rostilos.codecrow.security.web.jwt.AuthEntryPoint; import org.rostilos.codecrow.security.web.jwt.AuthTokenFilter; @@ -25,26 +26,23 @@ import org.springframework.web.cors.UrlBasedCorsConfigurationSource; @Configuration -@EnableMethodSecurity - (securedEnabled = true, - jsr250Enabled = true) +@EnableMethodSecurity(securedEnabled = true, jsr250Enabled = true) public class WebSecurityConfig { @Value("${codecrow.security.encryption-key}") private String encryptionKey; @Value("${codecrow.security.encryption-key-old}") private String oldEncryptionKey; - + @Value("${codecrow.internal.api.secret:}") private String internalApiSecret; - + private final UserDetailsServiceImpl userDetailsService; private final AuthEntryPoint unauthorizedHandler; public WebSecurityConfig( UserDetailsServiceImpl userDetailsService, - AuthEntryPoint unauthorizedHandler - ) { + AuthEntryPoint unauthorizedHandler) { this.userDetailsService = userDetailsService; this.unauthorizedHandler = unauthorizedHandler; } @@ -84,7 +82,8 @@ public CorsConfigurationSource corsConfigurationSource() { CorsConfiguration configuration = new CorsConfiguration(); // TODO: replace Arrays.asList("*") with the explicit domain(s) - // TODO: Example: Arrays.asList("http://localhost:8080", "https://frontend.rostilos.pp.ua") + // TODO: Example: Arrays.asList("http://localhost:8080", + // "https://frontend.rostilos.pp.ua") configuration.setAllowedOriginPatterns(Arrays.asList("*")); configuration.setAllowedMethods(Arrays.asList("GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS")); @@ -108,34 +107,37 @@ public SecurityFilterChain filterChain(HttpSecurity http) throws Exception { .headers(headers -> headers .frameOptions(HeadersConfigurer.FrameOptionsConfig::disable) .contentSecurityPolicy(csp -> csp - .policyDirectives("frame-ancestors 'self' https://bitbucket.org https://*.bitbucket.org") - ) - ) - .authorizeHttpRequests(auth -> - auth - // Allow all OPTIONS requests (CORS preflight) - .requestMatchers(org.springframework.http.HttpMethod.OPTIONS, "/**").permitAll() - .requestMatchers("/api/auth/**").permitAll() - .requestMatchers("/api/test/**").permitAll() - // OAuth callbacks need to be public (called by VCS providers) - .requestMatchers("/api/*/integrations/*/app/callback").permitAll() - // Generic OAuth callbacks without workspace slug (for GitHub, etc.) - .requestMatchers("/api/integrations/*/app/callback").permitAll() - .requestMatchers("/actuator/**").permitAll() - .requestMatchers("/internal/projects/**").permitAll() - .requestMatchers("/api/internal/**").permitAll() - .requestMatchers("/swagger-ui-custom.html").permitAll() - .requestMatchers("/api-docs").permitAll() - // Bitbucket Connect App lifecycle callbacks (uses JWT auth) - .requestMatchers("/api/bitbucket/connect/descriptor").permitAll() - .requestMatchers("/api/bitbucket/connect/installed").permitAll() - .requestMatchers("/api/bitbucket/connect/uninstalled").permitAll() - .requestMatchers("/api/bitbucket/connect/enabled").permitAll() - .requestMatchers("/api/bitbucket/connect/disabled").permitAll() - .requestMatchers("/api/bitbucket/connect/status").permitAll() - .requestMatchers("/api/bitbucket/connect/configure").permitAll() - .anyRequest().authenticated() - ); + .policyDirectives( + "frame-ancestors 'self' https://bitbucket.org https://*.bitbucket.org"))) + .authorizeHttpRequests(auth -> auth + // Allow async dispatches to complete (SSE, streaming responses) + .dispatcherTypeMatchers(DispatcherType.ASYNC, DispatcherType.ERROR).permitAll() + // Allow all OPTIONS requests (CORS preflight) + .requestMatchers(org.springframework.http.HttpMethod.OPTIONS, "/**").permitAll() + // Allow error page to be rendered without authentication + .requestMatchers("/error").permitAll() + .requestMatchers("/api/auth/**").permitAll() + .requestMatchers("/api/test/**").permitAll() + // OAuth callbacks need to be public (called by VCS providers) + .requestMatchers("/api/*/integrations/*/app/callback").permitAll() + // Generic OAuth callbacks without workspace slug (for GitHub, etc.) + .requestMatchers("/api/integrations/*/app/callback").permitAll() + .requestMatchers("/actuator/**").permitAll() + .requestMatchers("/internal/projects/**").permitAll() + .requestMatchers("/api/internal/**").permitAll() + .requestMatchers("/swagger-ui-custom.html").permitAll() + .requestMatchers("/api-docs").permitAll() + // Bitbucket Connect App lifecycle callbacks (uses JWT auth) + .requestMatchers("/api/bitbucket/connect/descriptor").permitAll() + .requestMatchers("/api/bitbucket/connect/installed").permitAll() + .requestMatchers("/api/bitbucket/connect/uninstalled").permitAll() + .requestMatchers("/api/bitbucket/connect/enabled").permitAll() + .requestMatchers("/api/bitbucket/connect/disabled").permitAll() + .requestMatchers("/api/bitbucket/connect/status").permitAll() + .requestMatchers("/api/bitbucket/connect/configure").permitAll() + // Stripe webhooks (authenticate via signature verification, not JWT) + .requestMatchers("/api/webhooks/**").permitAll() + .anyRequest().authenticated()); http.authenticationProvider(authenticationProvider()); diff --git a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClient.java b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClient.java index d6946c7d..ff7204d1 100644 --- a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClient.java +++ b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClient.java @@ -218,4 +218,39 @@ default List getRepositoryCollaborators(String workspaceId, Str default String getBranchDiff(String workspaceId, String repoIdOrSlug, String baseBranch, String compareBranch) throws IOException { throw new UnsupportedOperationException("Branch diff not supported by this provider"); } + + /** + * Batch fetch file contents from repository. + * Implementations should use provider-optimal strategies: + * - GitHub: GraphQL batch query + * - Bitbucket/GitLab: Parallel fetch with exponential backoff + * + * @param workspaceId the external workspace/org ID + * @param repoIdOrSlug the repository ID or slug + * @param filePaths list of file paths to fetch + * @param branchOrCommit branch name or commit hash + * @param maxFileSizeBytes maximum file size to fetch (skip larger files) + * @return map of path -> content (missing/skipped files not in map) + */ + default java.util.Map getFileContents( + String workspaceId, + String repoIdOrSlug, + List filePaths, + String branchOrCommit, + int maxFileSizeBytes + ) throws IOException { + // Default implementation: sequential fetch with size check + java.util.Map results = new java.util.HashMap<>(); + for (String path : filePaths) { + try { + String content = getFileContent(workspaceId, repoIdOrSlug, path, branchOrCommit); + if (content != null && content.length() <= maxFileSizeBytes) { + results.put(path, content); + } + } catch (IOException e) { + // Skip files that fail to fetch + } + } + return results; + } } diff --git a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClientProvider.java b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClientProvider.java index 1a6ad7de..7190f68b 100644 --- a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClientProvider.java +++ b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/VcsClientProvider.java @@ -210,7 +210,7 @@ public boolean needsTokenRefresh(VcsConnection connection) { * @return updated connection with new tokens * @throws VcsClientException if refresh fails */ - @Transactional + @Transactional(timeout = 30) // 30 second timeout to prevent deadlocks public VcsConnection refreshToken(VcsConnection connection) { log.info("Refreshing access token for connection: {} (provider: {}, type: {})", connection.getId(), connection.getProviderType(), connection.getConnectionType()); @@ -415,7 +415,12 @@ private TokenResponse refreshGitLabToken(String refreshToken) throws IOException throw new IOException("GitLab OAuth credentials not configured. Set codecrow.gitlab.oauth.client-id and codecrow.gitlab.oauth.client-secret"); } - OkHttpClient httpClient = new OkHttpClient(); + // Use short timeouts to prevent holding database locks during slow network operations + OkHttpClient httpClient = new OkHttpClient.Builder() + .connectTimeout(10, java.util.concurrent.TimeUnit.SECONDS) + .readTimeout(15, java.util.concurrent.TimeUnit.SECONDS) + .writeTimeout(10, java.util.concurrent.TimeUnit.SECONDS) + .build(); // Determine GitLab token URL (support self-hosted) String tokenUrl = (gitlabBaseUrl != null && !gitlabBaseUrl.isBlank() && !gitlabBaseUrl.equals("https://gitlab.com")) diff --git a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/BitbucketCloudClient.java b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/BitbucketCloudClient.java index 9c3e4b46..41989b9f 100644 --- a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/BitbucketCloudClient.java +++ b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/BitbucketCloudClient.java @@ -5,6 +5,8 @@ import okhttp3.*; import org.rostilos.codecrow.vcsclient.VcsClient; import org.rostilos.codecrow.vcsclient.model.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.IOException; import java.io.InputStream; @@ -20,6 +22,7 @@ */ public class BitbucketCloudClient implements VcsClient { + private static final Logger log = LoggerFactory.getLogger(BitbucketCloudClient.class); private static final String API_BASE = "https://api.bitbucket.org/2.0"; private static final int DEFAULT_PAGE_SIZE = 50; private static final MediaType JSON_MEDIA_TYPE = MediaType.parse("application/json"); @@ -791,4 +794,66 @@ private VcsCollaborator parseCollaboratorPermission(JsonNode permNode) { return new VcsCollaborator(userId, username, displayName, avatarUrl, permission, htmlUrl); } + + /** + * Batch fetch file contents with parallel execution and exponential backoff. + * Bitbucket Cloud doesn't have a batch API, so we fetch in parallel with rate limit handling. + */ + @Override + public java.util.Map getFileContents( + String workspaceId, + String repoIdOrSlug, + java.util.List filePaths, + String branchOrCommit, + int maxFileSizeBytes + ) throws IOException { + java.util.Map results = new java.util.concurrent.ConcurrentHashMap<>(); + + // Use parallel stream with controlled concurrency + int parallelism = Math.min(10, filePaths.size()); // Max 10 concurrent requests + java.util.concurrent.ForkJoinPool customPool = new java.util.concurrent.ForkJoinPool(parallelism); + + try { + customPool.submit(() -> + filePaths.parallelStream().forEach(path -> { + int maxRetries = 3; + int retryCount = 0; + long backoffMs = 1000; // Start with 1 second + + while (retryCount < maxRetries) { + try { + String content = getFileContent(workspaceId, repoIdOrSlug, path, branchOrCommit); + if (content != null && content.getBytes(java.nio.charset.StandardCharsets.UTF_8).length <= maxFileSizeBytes) { + results.put(path, content); + } + break; // Success, exit retry loop + } catch (IOException e) { + retryCount++; + if (e.getMessage() != null && e.getMessage().contains("429")) { + // Rate limited - exponential backoff + try { + Thread.sleep(backoffMs); + backoffMs *= 2; // Double the backoff + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + break; + } + } else if (retryCount >= maxRetries) { + // Log and skip this file + log.warn("Failed to fetch file {} after {} retries: {}", path, maxRetries, e.getMessage()); + } + } + } + }) + ).get(); + } catch (InterruptedException | java.util.concurrent.ExecutionException e) { + log.error("Error in parallel file fetch: {}", e.getMessage()); + throw new IOException("Batch file fetch failed", e); + } finally { + customPool.shutdown(); + } + + log.info("Batch fetched {}/{} files from Bitbucket", results.size(), filePaths.size()); + return results; + } } diff --git a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/actions/CheckFileExistsInBranchAction.java b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/actions/CheckFileExistsInBranchAction.java index 143c017b..c61a5e93 100644 --- a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/actions/CheckFileExistsInBranchAction.java +++ b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/bitbucket/cloud/actions/CheckFileExistsInBranchAction.java @@ -15,10 +15,14 @@ /** * Action to check if a file exists in a specific branch on Bitbucket Cloud. * Uses the Bitbucket Cloud API to verify file existence without downloading content. + * Includes retry logic with exponential backoff for rate-limited (429) responses. */ public class CheckFileExistsInBranchAction { private static final Logger log = LoggerFactory.getLogger(CheckFileExistsInBranchAction.class); + private static final int MAX_RETRIES = 3; + private static final long INITIAL_BACKOFF_MS = 2_000; + private final OkHttpClient authorizedOkHttpClient; public CheckFileExistsInBranchAction(OkHttpClient authorizedOkHttpClient) { @@ -28,18 +32,19 @@ public CheckFileExistsInBranchAction(OkHttpClient authorizedOkHttpClient) { /** * Checks if a file exists in the specified branch. * Uses HEAD request to check existence without downloading file content. + * Retries with exponential backoff on 429 (rate-limit) responses. * * @param workspace workspace or team slug * @param repoSlug repository slug * @param branchName branch name (or commit hash) * @param filePath file path relative to repository root * @return true if file exists in the branch, false otherwise - * @throws IOException on network errors + * @throws IOException on network errors after retries are exhausted */ public boolean fileExists(String workspace, String repoSlug, String branchName, String filePath) throws IOException { String ws = Optional.ofNullable(workspace).orElse(""); String encodedPath = encodeFilePath(filePath); - + // Use Bitbucket Cloud API endpoint to check file existence with HEAD request String apiUrl = String.format("%s/repositories/%s/%s/src/%s/%s", BitbucketCloudConfig.BITBUCKET_API_BASE, ws, repoSlug, branchName, encodedPath); @@ -49,22 +54,53 @@ public boolean fileExists(String workspace, String repoSlug, String branchName, .head() .build(); - try (Response resp = authorizedOkHttpClient.newCall(req).execute()) { - if (resp.isSuccessful()) { - return true; - } else if (resp.code() == 404) { - log.debug("File not found: {} in branch {} (workspace: {}, repo: {})", - filePath, branchName, workspace, repoSlug); - return false; - } else { - String msg = String.format("Unexpected response %d when checking file existence: %s in branch %s", - resp.code(), filePath, branchName); - log.warn(msg); - throw new IOException(msg); + int attempt = 0; + long backoffMs = INITIAL_BACKOFF_MS; + + while (true) { + try (Response resp = authorizedOkHttpClient.newCall(req).execute()) { + if (resp.isSuccessful()) { + return true; + } else if (resp.code() == 404) { + log.debug("File not found: {} in branch {} (workspace: {}, repo: {})", + filePath, branchName, workspace, repoSlug); + return false; + } else if (resp.code() == 429 && attempt < MAX_RETRIES) { + // Rate limited — honour Retry-After header if present, otherwise exponential backoff + long waitMs = backoffMs; + String retryAfter = resp.header("Retry-After"); + if (retryAfter != null) { + try { + waitMs = Long.parseLong(retryAfter.trim()) * 1_000; + } catch (NumberFormatException ignored) { + // Use default backoff + } + } + log.info("Rate limited (429) checking file {}. Retrying in {}ms (attempt {}/{})", + filePath, waitMs, attempt + 1, MAX_RETRIES); + try { + Thread.sleep(waitMs); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while waiting for rate-limit backoff", ie); + } + attempt++; + backoffMs *= 2; // Exponential backoff + } else { + String msg = String.format("Unexpected response %d when checking file existence: %s in branch %s", + resp.code(), filePath, branchName); + log.warn(msg); + throw new IOException(msg); + } + } catch (IOException e) { + if (attempt < MAX_RETRIES && e.getMessage() != null && e.getMessage().contains("429")) { + attempt++; + backoffMs *= 2; + continue; + } + log.error("Failed to check file existence for {}: {}", filePath, e.getMessage(), e); + throw e; } - } catch (IOException e) { - log.error("Failed to check file existence for {}: {}", filePath, e.getMessage(), e); - throw e; } } diff --git a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/github/GitHubClient.java b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/github/GitHubClient.java index d9a6de5e..e36ce4a4 100644 --- a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/github/GitHubClient.java +++ b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/github/GitHubClient.java @@ -933,6 +933,115 @@ private Request createGetRequest(String url) { .build(); } + /** + * Batch fetch file contents using GitHub GraphQL API for efficiency. + * GraphQL allows fetching multiple files in a single request. + */ + @Override + public java.util.Map getFileContents( + String workspaceId, + String repoIdOrSlug, + java.util.List filePaths, + String branchOrCommit, + int maxFileSizeBytes + ) throws IOException { + java.util.Map results = new java.util.HashMap<>(); + + // GitHub GraphQL API endpoint + String graphqlUrl = "https://api.github.com/graphql"; + + // Process in batches of 50 files (GraphQL complexity limits) + int batchSize = 50; + for (int i = 0; i < filePaths.size(); i += batchSize) { + java.util.List batch = filePaths.subList(i, Math.min(i + batchSize, filePaths.size())); + + // Build GraphQL query for this batch + StringBuilder queryBuilder = new StringBuilder(); + queryBuilder.append("query { repository(owner: \"").append(workspaceId) + .append("\", name: \"").append(repoIdOrSlug).append("\") {"); + + for (int j = 0; j < batch.size(); j++) { + String path = batch.get(j); + String alias = "file" + j; + String expression = branchOrCommit + ":" + path; + queryBuilder.append(alias).append(": object(expression: \"") + .append(expression.replace("\"", "\\\"")) + .append("\") { ... on Blob { text byteSize } } "); + } + queryBuilder.append("}}"); + + String query = queryBuilder.toString(); + String requestBody = objectMapper.writeValueAsString(java.util.Map.of("query", query)); + + Request request = new Request.Builder() + .url(graphqlUrl) + .header(ACCEPT_HEADER, GITHUB_ACCEPT_HEADER) + .header(GITHUB_API_VERSION_HEADER, GITHUB_API_VERSION) + .post(RequestBody.create(requestBody, JSON_MEDIA_TYPE)) + .build(); + + int maxRetries = 3; + int retryCount = 0; + long backoffMs = 1000; + + while (retryCount < maxRetries) { + try (Response response = httpClient.newCall(request).execute()) { + if (response.code() == 429) { + // Rate limited - exponential backoff + retryCount++; + Thread.sleep(backoffMs); + backoffMs *= 2; + continue; + } + + if (!response.isSuccessful()) { + log.warn("GraphQL batch fetch failed with status {}, falling back to REST API", response.code()); + // Fall back to sequential REST API + for (String path : batch) { + try { + String content = getFileContent(workspaceId, repoIdOrSlug, path, branchOrCommit); + if (content != null && content.getBytes(java.nio.charset.StandardCharsets.UTF_8).length <= maxFileSizeBytes) { + results.put(path, content); + } + } catch (IOException e) { + log.debug("Skipping file {}: {}", path, e.getMessage()); + } + } + break; + } + + // Parse GraphQL response + JsonNode root = objectMapper.readTree(response.body().string()); + JsonNode data = root.path("data").path("repository"); + + for (int j = 0; j < batch.size(); j++) { + String path = batch.get(j); + String alias = "file" + j; + JsonNode fileNode = data.path(alias); + + if (!fileNode.isMissingNode() && fileNode.has("text")) { + int byteSize = fileNode.path("byteSize").asInt(0); + if (byteSize <= maxFileSizeBytes) { + String text = fileNode.get("text").asText(); + if (text != null) { + results.put(path, text); + } + } + } + } + break; // Success + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted during batch fetch", e); + } + } + } + + log.info("Batch fetched {}/{} files from GitHub via GraphQL", results.size(), filePaths.size()); + return results; + } + private record GitHubWebhookRequest( String name, diff --git a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/gitlab/GitLabClient.java b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/gitlab/GitLabClient.java index 7e50f566..0fb1ec93 100644 --- a/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/gitlab/GitLabClient.java +++ b/java-ecosystem/libs/vcs-client/src/main/java/org/rostilos/codecrow/vcsclient/gitlab/GitLabClient.java @@ -498,6 +498,77 @@ public String getLatestCommitHash(String workspaceId, String repoIdOrSlug, Strin } } + @Override + public String getBranchDiff(String workspaceId, String repoIdOrSlug, String baseBranch, String compareBranch) throws IOException { + // GitLab: GET /projects/:id/repository/compare + // Returns diff between two branches/commits + // API: https://docs.gitlab.com/ee/api/repositories.html#compare-branches-tags-or-commits + String projectPath = workspaceId + "/" + repoIdOrSlug; + String encodedPath = URLEncoder.encode(projectPath, StandardCharsets.UTF_8); + String encodedFrom = URLEncoder.encode(baseBranch, StandardCharsets.UTF_8); + String encodedTo = URLEncoder.encode(compareBranch, StandardCharsets.UTF_8); + + String url = baseUrl + "/projects/" + encodedPath + "/repository/compare?from=" + encodedFrom + "&to=" + encodedTo; + + Request request = createGetRequest(url); + try (Response response = httpClient.newCall(request).execute()) { + if (!response.isSuccessful()) { + throw createException("get branch diff", response); + } + + JsonNode root = objectMapper.readTree(response.body().string()); + JsonNode diffs = root.get("diffs"); + + if (diffs == null || !diffs.isArray() || diffs.isEmpty()) { + return ""; + } + + // Build unified diff format from GitLab's compare response + StringBuilder diffBuilder = new StringBuilder(); + for (JsonNode diff : diffs) { + String oldPath = getTextOrNull(diff, "old_path"); + String newPath = getTextOrNull(diff, "new_path"); + boolean newFile = diff.has("new_file") && diff.get("new_file").asBoolean(); + boolean deletedFile = diff.has("deleted_file") && diff.get("deleted_file").asBoolean(); + boolean renamedFile = diff.has("renamed_file") && diff.get("renamed_file").asBoolean(); + String diffContent = getTextOrNull(diff, "diff"); + + // Build git diff header + diffBuilder.append("diff --git a/").append(oldPath).append(" b/").append(newPath).append("\n"); + + if (newFile) { + diffBuilder.append("new file mode 100644\n"); + } else if (deletedFile) { + diffBuilder.append("deleted file mode 100644\n"); + } else if (renamedFile) { + diffBuilder.append("rename from ").append(oldPath).append("\n"); + diffBuilder.append("rename to ").append(newPath).append("\n"); + } + + // Proper unified diff headers: /dev/null for new/deleted files + if (newFile) { + diffBuilder.append("--- /dev/null\n"); + diffBuilder.append("+++ b/").append(newPath).append("\n"); + } else if (deletedFile) { + diffBuilder.append("--- a/").append(oldPath).append("\n"); + diffBuilder.append("+++ /dev/null\n"); + } else { + diffBuilder.append("--- a/").append(oldPath).append("\n"); + diffBuilder.append("+++ b/").append(newPath).append("\n"); + } + + if (diffContent != null && !diffContent.isEmpty()) { + diffBuilder.append(diffContent); + if (!diffContent.endsWith("\n")) { + diffBuilder.append("\n"); + } + } + } + + return diffBuilder.toString(); + } + } + @Override public List listBranches(String workspaceId, String repoIdOrSlug) throws IOException { List branches = new ArrayList<>(); @@ -778,4 +849,66 @@ private Request createGetRequest(String url) { .get() .build(); } + + /** + * Batch fetch file contents with parallel execution and exponential backoff. + * GitLab doesn't have a batch API, so we fetch in parallel with rate limit handling. + */ + @Override + public java.util.Map getFileContents( + String workspaceId, + String repoIdOrSlug, + java.util.List filePaths, + String branchOrCommit, + int maxFileSizeBytes + ) throws IOException { + java.util.Map results = new java.util.concurrent.ConcurrentHashMap<>(); + + // Use parallel stream with controlled concurrency + int parallelism = Math.min(10, filePaths.size()); // Max 10 concurrent requests + java.util.concurrent.ForkJoinPool customPool = new java.util.concurrent.ForkJoinPool(parallelism); + + try { + customPool.submit(() -> + filePaths.parallelStream().forEach(path -> { + int maxRetries = 3; + int retryCount = 0; + long backoffMs = 1000; // Start with 1 second + + while (retryCount < maxRetries) { + try { + String content = getFileContent(workspaceId, repoIdOrSlug, path, branchOrCommit); + if (content != null && content.getBytes(java.nio.charset.StandardCharsets.UTF_8).length <= maxFileSizeBytes) { + results.put(path, content); + } + break; // Success, exit retry loop + } catch (IOException e) { + retryCount++; + if (e.getMessage() != null && (e.getMessage().contains("429") || e.getMessage().contains("rate limit"))) { + // Rate limited - exponential backoff + try { + Thread.sleep(backoffMs); + backoffMs *= 2; // Double the backoff + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + break; + } + } else if (retryCount >= maxRetries) { + // Log and skip this file + log.warn("Failed to fetch file {} after {} retries: {}", path, maxRetries, e.getMessage()); + } + } + } + }) + ).get(); + } catch (InterruptedException | java.util.concurrent.ExecutionException e) { + log.error("Error in parallel file fetch: {}", e.getMessage()); + throw new IOException("Batch file fetch failed", e); + } finally { + customPool.shutdown(); + } + + log.info("Batch fetched {}/{} files from GitLab", results.size(), filePaths.size()); + return results; + } } diff --git a/java-ecosystem/mcp-servers/platform-mcp/src/main/java/org/rostilos/codecrow/platformmcp/db/PlatformDbService.java b/java-ecosystem/mcp-servers/platform-mcp/src/main/java/org/rostilos/codecrow/platformmcp/db/PlatformDbService.java deleted file mode 100644 index e69de29b..00000000 diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/module-info.java b/java-ecosystem/services/pipeline-agent/src/main/java/module-info.java index 739f411f..3188b03a 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/module-info.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/module-info.java @@ -9,6 +9,7 @@ requires spring.beans; requires org.slf4j; requires jakarta.validation; + requires jakarta.persistence; requires spring.web; requires jjwt.api; requires okhttp3; diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/service/BitbucketAiClientService.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/service/BitbucketAiClientService.java index 0db3cee3..7060f11e 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/service/BitbucketAiClientService.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/service/BitbucketAiClientService.java @@ -1,6 +1,8 @@ package org.rostilos.codecrow.pipelineagent.bitbucket.service; import okhttp3.OkHttpClient; +import org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment.PrEnrichmentDataDto; +import org.rostilos.codecrow.analysisengine.service.PrFileEnrichmentService; import org.rostilos.codecrow.core.model.ai.AIConnection; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisMode; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisType; @@ -13,10 +15,13 @@ import org.rostilos.codecrow.analysisengine.dto.request.processor.BranchProcessRequest; import org.rostilos.codecrow.analysisengine.dto.request.processor.PrProcessRequest; import org.rostilos.codecrow.analysisengine.dto.request.ai.AiAnalysisRequest; +import org.rostilos.codecrow.analysisengine.exception.DiffTooLargeException; import org.rostilos.codecrow.analysisengine.service.vcs.VcsAiClientService; import org.rostilos.codecrow.analysisengine.util.DiffContentFilter; import org.rostilos.codecrow.analysisengine.util.DiffParser; +import org.rostilos.codecrow.analysisengine.util.TokenEstimator; import org.rostilos.codecrow.security.oauth.TokenEncryptionService; +import org.rostilos.codecrow.vcsclient.VcsClient; import org.rostilos.codecrow.vcsclient.VcsClientProvider; import org.rostilos.codecrow.vcsclient.bitbucket.cloud.actions.GetCommitRangeDiffAction; import org.rostilos.codecrow.vcsclient.bitbucket.cloud.actions.GetPullRequestAction; @@ -25,6 +30,7 @@ import org.rostilos.codecrow.vcsclient.utils.VcsConnectionCredentialsExtractor.VcsConnectionCredentials; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Service; import java.io.IOException; @@ -52,14 +58,17 @@ public class BitbucketAiClientService implements VcsAiClientService { private final TokenEncryptionService tokenEncryptionService; private final VcsClientProvider vcsClientProvider; private final VcsConnectionCredentialsExtractor credentialsExtractor; + private final PrFileEnrichmentService enrichmentService; public BitbucketAiClientService( TokenEncryptionService tokenEncryptionService, - VcsClientProvider vcsClientProvider + VcsClientProvider vcsClientProvider, + @Autowired(required = false) PrFileEnrichmentService enrichmentService ) { this.tokenEncryptionService = tokenEncryptionService; this.vcsClientProvider = vcsClientProvider; this.credentialsExtractor = new VcsConnectionCredentialsExtractor(tokenEncryptionService); + this.enrichmentService = enrichmentService; } @Override @@ -94,19 +103,37 @@ public AiAnalysisRequest buildAiAnalysisRequest( if(request.getAnalysisType() == AnalysisType.BRANCH_ANALYSIS){ return buildBranchAnalysisRequest(project, (BranchProcessRequest) request, previousAnalysis); } else { - return buildPrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis); + return buildPrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis, Collections.emptyList()); + } + } + + @Override + public AiAnalysisRequest buildAiAnalysisRequest( + Project project, + AnalysisProcessRequest request, + Optional previousAnalysis, + List allPrAnalyses + ) throws GeneralSecurityException { + if(request.getAnalysisType() == AnalysisType.BRANCH_ANALYSIS){ + return buildBranchAnalysisRequest(project, (BranchProcessRequest) request, previousAnalysis); + } else { + return buildPrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis, allPrAnalyses); } } public AiAnalysisRequest buildPrAnalysisRequest( Project project, PrProcessRequest request, - Optional previousAnalysis + Optional previousAnalysis, + List allPrAnalyses ) throws GeneralSecurityException { VcsInfo vcsInfo = getVcsInfo(project); VcsConnection vcsConnection = vcsInfo.vcsConnection(); AIConnection aiConnection = project.getAiBinding().getAiConnection(); - AIConnection projectAiConnection = project.getAiBinding().getAiConnection(); + + // CRITICAL: Log the AI connection being used for debugging + log.info("Building PR analysis request for project={}, AI model={}, provider={}, aiConnectionId={}", + project.getId(), aiConnection.getAiModel(), aiConnection.getProviderKey(), aiConnection.getId()); // Initialize variables List changedFiles = Collections.emptyList(); @@ -157,6 +184,23 @@ public AiAnalysisRequest buildPrAnalysisRequest( originalSize > 0 ? (100 - (filteredSize * 100 / originalSize)) : 0); } + // Check token limit before proceeding with analysis + int maxTokenLimit = project.getEffectiveConfig().maxAnalysisTokenLimit(); + TokenEstimator.TokenEstimationResult tokenEstimate = TokenEstimator.estimateAndCheck(rawDiff, maxTokenLimit); + log.info("Token estimation for PR diff: {}", tokenEstimate.toLogString()); + + if (tokenEstimate.exceedsLimit()) { + log.warn("PR diff exceeds token limit - skipping analysis. Project={}, PR={}, Tokens={}/{}", + project.getId(), request.getPullRequestId(), + tokenEstimate.estimatedTokens(), tokenEstimate.maxAllowedTokens()); + throw new DiffTooLargeException( + tokenEstimate.estimatedTokens(), + tokenEstimate.maxAllowedTokens(), + project.getId(), + request.getPullRequestId() + ); + } + // Determine analysis mode: INCREMENTAL if we have previous analysis with different commit boolean canUseIncremental = previousAnalysis.isPresent() && previousCommitHash != null @@ -209,15 +253,35 @@ public AiAnalysisRequest buildPrAnalysisRequest( // Continue without metadata - RAG will use fallback } - var builder = AiAnalysisRequestImpl.builder() + // Enrich PR with full file contents and dependency graph + PrEnrichmentDataDto enrichmentData = PrEnrichmentDataDto.empty(); + if (enrichmentService != null && enrichmentService.isEnrichmentEnabled() && !changedFiles.isEmpty()) { + try { + VcsClient vcsClient = vcsClientProvider.getClient(vcsConnection); + enrichmentData = enrichmentService.enrichPrFiles( + vcsClient, + vcsInfo.workspace(), + vcsInfo.repoSlug(), + request.getSourceBranchName(), + changedFiles + ); + log.info("PR enrichment completed: {} files enriched, {} relationships", + enrichmentData.stats().filesEnriched(), + enrichmentData.stats().relationshipsFound()); + } catch (Exception e) { + log.warn("Failed to enrich PR files (non-critical): {}", e.getMessage()); + } + } + + AiAnalysisRequestImpl.Builder builder = AiAnalysisRequestImpl.builder() .withProjectId(project.getId()) .withPullRequestId(request.getPullRequestId()) .withProjectAiConnection(aiConnection) .withProjectVcsConnectionBindingInfo(vcsInfo.workspace(), vcsInfo.repoSlug()) - .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(projectAiConnection.getApiKeyEncrypted())) + .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted())) .withUseLocalMcp(true) - .withPreviousAnalysisData(previousAnalysis) - .withMaxAllowedTokens(aiConnection.getTokenLimitation()) + .withAllPrAnalysesData(allPrAnalyses) // Use full PR history instead of just previous version + .withMaxAllowedTokens(project.getEffectiveConfig().maxAnalysisTokenLimit()) .withAnalysisType(request.getAnalysisType()) .withPrTitle(prTitle) .withPrDescription(prDescription) @@ -231,7 +295,9 @@ public AiAnalysisRequest buildPrAnalysisRequest( .withAnalysisMode(analysisMode) .withDeltaDiff(deltaDiff) .withPreviousCommitHash(previousCommitHash) - .withCurrentCommitHash(currentCommitHash); + .withCurrentCommitHash(currentCommitHash) + // File enrichment data + .withEnrichmentData(enrichmentData); // Add VCS credentials based on connection type addVcsCredentials(builder, vcsConnection); @@ -278,17 +344,16 @@ public AiAnalysisRequest buildBranchAnalysisRequest( VcsInfo vcsInfo = getVcsInfo(project); VcsConnection vcsConnection = vcsInfo.vcsConnection(); AIConnection aiConnection = project.getAiBinding().getAiConnection(); - AIConnection projectAiConnection = project.getAiBinding().getAiConnection(); var builder = AiAnalysisRequestImpl.builder() .withProjectId(project.getId()) .withPullRequestId(null) .withProjectAiConnection(aiConnection) .withProjectVcsConnectionBindingInfo(vcsInfo.workspace(), vcsInfo.repoSlug()) - .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(projectAiConnection.getApiKeyEncrypted())) + .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted())) .withUseLocalMcp(true) .withPreviousAnalysisData(previousAnalysis) - .withMaxAllowedTokens(aiConnection.getTokenLimitation()) + .withMaxAllowedTokens(project.getEffectiveConfig().maxAnalysisTokenLimit()) .withAnalysisType(request.getAnalysisType()) .withTargetBranchName(request.getTargetBranchName()) .withCurrentCommitHash(request.getCommitHash()) diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/webhookhandler/BitbucketCloudBranchWebhookHandler.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/webhookhandler/BitbucketCloudBranchWebhookHandler.java index 81dd08c9..4b6a4a9e 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/webhookhandler/BitbucketCloudBranchWebhookHandler.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/bitbucket/webhookhandler/BitbucketCloudBranchWebhookHandler.java @@ -7,6 +7,7 @@ import org.rostilos.codecrow.analysisengine.dto.request.processor.BranchProcessRequest; import org.rostilos.codecrow.analysisengine.processor.analysis.BranchAnalysisProcessor; import org.rostilos.codecrow.pipelineagent.generic.dto.webhook.WebhookPayload; +import org.rostilos.codecrow.pipelineagent.generic.service.WebhookDeduplicationService; import org.rostilos.codecrow.pipelineagent.generic.webhookhandler.AbstractWebhookHandler; import org.rostilos.codecrow.pipelineagent.generic.webhookhandler.WebhookHandler; import org.slf4j.Logger; @@ -38,13 +39,16 @@ public class BitbucketCloudBranchWebhookHandler extends AbstractWebhookHandler i private final BranchAnalysisProcessor branchAnalysisProcessor; private final RagOperationsService ragOperationsService; + private final WebhookDeduplicationService deduplicationService; public BitbucketCloudBranchWebhookHandler( BranchAnalysisProcessor branchAnalysisProcessor, - @Autowired(required = false) RagOperationsService ragOperationsService + @Autowired(required = false) RagOperationsService ragOperationsService, + WebhookDeduplicationService deduplicationService ) { this.branchAnalysisProcessor = branchAnalysisProcessor; this.ragOperationsService = ragOperationsService; + this.deduplicationService = deduplicationService; } @Override @@ -87,6 +91,12 @@ public WebhookResult handle(WebhookPayload payload, Project project, Consumer> eventConsumer) { String placeholderCommentId = null; + String acquiredLockKey = null; try { + // Try to acquire lock atomically BEFORE posting placeholder + // This prevents race condition where multiple webhooks could post duplicate placeholders + String sourceBranch = payload.sourceBranch(); + Optional earlyLock = analysisLockService.acquireLock( + project, sourceBranch, AnalysisLockType.PR_ANALYSIS, + payload.commitHash(), Long.parseLong(payload.pullRequestId())); + + if (earlyLock.isEmpty()) { + log.info("PR analysis already in progress for project={}, branch={}, PR={} - skipping duplicate webhook", + project.getId(), sourceBranch, payload.pullRequestId()); + return WebhookResult.ignored("PR analysis already in progress for this branch"); + } + + acquiredLockKey = earlyLock.get(); + + // Lock acquired - placeholder posting is now protected from race conditions + // Post placeholder comment immediately to show analysis has started placeholderCommentId = postPlaceholderComment(project, Long.parseLong(payload.pullRequestId())); @@ -118,6 +144,8 @@ private WebhookResult handlePullRequestEvent(WebhookPayload payload, Project pro request.placeholderCommentId = placeholderCommentId; request.prAuthorId = payload.prAuthorId(); request.prAuthorUsername = payload.prAuthorUsername(); + // Pass the pre-acquired lock key to avoid double-locking in the processor + request.preAcquiredLockKey = acquiredLockKey; log.info("Processing PR analysis: project={}, PR={}, source={}, target={}, placeholderCommentId={}", project.getId(), request.pullRequestId, request.sourceBranchName, request.targetBranchName, placeholderCommentId); @@ -143,8 +171,20 @@ private WebhookResult handlePullRequestEvent(WebhookPayload payload, Project pro return WebhookResult.success("PR analysis completed", result); + } catch (DiffTooLargeException | AnalysisLockedException e) { + // Re-throw these exceptions so WebhookAsyncProcessor can handle them properly + // Release the lock since processor won't take ownership + if (acquiredLockKey != null) { + analysisLockService.releaseLock(acquiredLockKey); + } + log.warn("PR analysis failed with recoverable exception for project {}: {}", project.getId(), e.getMessage()); + throw e; } catch (Exception e) { log.error("PR analysis failed for project {}", project.getId(), e); + // Release the lock since processor won't take ownership + if (acquiredLockKey != null) { + analysisLockService.releaseLock(acquiredLockKey); + } // Try to update placeholder with error message if (placeholderCommentId != null) { try { diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/config/AsyncConfig.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/config/AsyncConfig.java index 0f119827..b70781db 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/config/AsyncConfig.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/config/AsyncConfig.java @@ -46,12 +46,21 @@ public Executor webhookExecutor() { ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor(); executor.setCorePoolSize(4); executor.setMaxPoolSize(8); - executor.setQueueCapacity(100); + executor.setQueueCapacity(200); // Increased to handle bursts executor.setThreadNamePrefix("webhook-"); executor.setWaitForTasksToCompleteOnShutdown(true); executor.setAwaitTerminationSeconds(120); + // Use AbortPolicy - throw RejectedExecutionException instead of blocking caller + // This allows webhook endpoints to return HTTP 503 and trigger retry from VCS + // Running in caller thread would block the HTTP response and cause timeout + executor.setRejectedExecutionHandler((r, e) -> { + log.error("WEBHOOK EXECUTOR REJECTED TASK! Queue is full. Pool size: {}, Active: {}, Queue size: {}", + e.getPoolSize(), e.getActiveCount(), e.getQueue().size()); + throw new java.util.concurrent.RejectedExecutionException( + "Webhook executor queue is full - cannot accept new tasks"); + }); executor.initialize(); - log.info("Webhook executor initialized with core={}, max={}", 4, 8); + log.info("Webhook executor initialized with core={}, max={}, queueCapacity={}", 4, 8, 200); return executor; } diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/controller/ProviderWebhookController.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/controller/ProviderWebhookController.java index d475014f..95ecfc0a 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/controller/ProviderWebhookController.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/controller/ProviderWebhookController.java @@ -260,6 +260,9 @@ private ResponseEntity processWebhook(EVcsProvider provider, WebhookPayload p String jobUrl = buildJobUrl(project, job); String logsStreamUrl = buildJobLogsStreamUrl(job); + log.info("Dispatching webhook to async processor: job={}, event={}", + job.getExternalId(), payload.eventType()); + // Process webhook asynchronously with proper transactional context webhookAsyncProcessor.processWebhookAsync( provider, @@ -269,6 +272,8 @@ private ResponseEntity processWebhook(EVcsProvider provider, WebhookPayload p job ); + log.info("Webhook dispatched to async processor: job={}", job.getExternalId()); + return ResponseEntity.accepted().body(Map.of( "status", "accepted", "message", "Webhook received, processing started", diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/WebhookAsyncProcessor.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/WebhookAsyncProcessor.java index a216316a..a1c67990 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/WebhookAsyncProcessor.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/WebhookAsyncProcessor.java @@ -6,69 +6,43 @@ import org.rostilos.codecrow.core.persistence.repository.project.ProjectRepository; import org.rostilos.codecrow.core.service.JobService; import org.rostilos.codecrow.pipelineagent.generic.dto.webhook.WebhookPayload; +import org.rostilos.codecrow.pipelineagent.generic.utils.CommentPlaceholders; import org.rostilos.codecrow.pipelineagent.generic.webhookhandler.WebhookHandler; +import org.rostilos.codecrow.analysisengine.exception.AnalysisLockedException; +import org.rostilos.codecrow.analysisengine.exception.DiffTooLargeException; import org.rostilos.codecrow.analysisengine.service.vcs.VcsReportingService; import org.rostilos.codecrow.analysisengine.service.vcs.VcsServiceFactory; import org.slf4j.Logger; import java.io.IOException; import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.context.annotation.Lazy; import org.springframework.scheduling.annotation.Async; import org.springframework.stereotype.Service; import org.springframework.transaction.annotation.Transactional; /** - * Service for processing webhooks asynchronously within a transactional context. - * This ensures Hibernate sessions are available for lazy loading. + * Service for processing webhooks asynchronously. + * + * Uses a short read-only transaction to load the project entity with all lazy associations, + * then processes the webhook (including the potentially long-running AI analysis) WITHOUT + * holding a DB transaction open. All downstream services manage their own transactions. */ @Service public class WebhookAsyncProcessor { private static final Logger log = LoggerFactory.getLogger(WebhookAsyncProcessor.class); - /** Comment markers for CodeCrow command responses */ - private static final String CODECROW_COMMAND_MARKER = ""; - private static final String CODECROW_SUMMARY_MARKER = ""; - private static final String CODECROW_REVIEW_MARKER = ""; - - /** Placeholder messages for commands */ - private static final String PLACEHOLDER_ANALYZE = """ - 🔄 **CodeCrow is analyzing this PR...** - - This may take a few minutes depending on the size of the changes. - I'll update this comment with the results when the analysis is complete. - """; - - private static final String PLACEHOLDER_SUMMARIZE = """ - 🔄 **CodeCrow is generating a summary...** - - I'm analyzing the changes and creating diagrams. - This comment will be updated with the summary when ready. - """; - - private static final String PLACEHOLDER_REVIEW = """ - 🔄 **CodeCrow is reviewing this PR...** - - I'm examining the code changes for potential issues. - This comment will be updated with the review results when complete. - """; - - private static final String PLACEHOLDER_ASK = """ - 🔄 **CodeCrow is processing your question...** - - I'm analyzing the context to provide a helpful answer. - """; - - private static final String PLACEHOLDER_DEFAULT = """ - 🔄 **CodeCrow is processing...** - - Please wait while I complete this task. - """; - private final ProjectRepository projectRepository; private final JobService jobService; private final VcsServiceFactory vcsServiceFactory; + // Self-injection for @Transactional proxy to work from @Async method + @Autowired + @Lazy + private WebhookAsyncProcessor self; + public WebhookAsyncProcessor( ProjectRepository projectRepository, JobService jobService, @@ -80,10 +54,11 @@ public WebhookAsyncProcessor( } /** - * Process a webhook asynchronously with proper transactional context. + * Process a webhook asynchronously. + * Delegates to processWebhookInTransaction via self-injection to ensure + * the @Transactional proxy on loadAndInitializeProject is invoked properly. */ @Async("webhookExecutor") - @Transactional public void processWebhookAsync( EVcsProvider provider, Long projectId, @@ -91,18 +66,70 @@ public void processWebhookAsync( WebhookHandler handler, Job job ) { + log.info("processWebhookAsync started for job {} (projectId={}, event={})", + job.getExternalId(), projectId, payload.eventType()); + try { + // Delegate to transactional method via self-reference to ensure proxy is used + self.processWebhookInTransaction(provider, projectId, payload, handler, job); + log.info("processWebhookAsync completed normally for job {}", job.getExternalId()); + } catch (Exception e) { + log.error("processWebhookAsync FAILED for job {}: {}", job.getExternalId(), e.getMessage(), e); + // Try to fail the job so it doesn't stay in PENDING + try { + log.info("Marking job {} as FAILED from async wrapper", job.getExternalId()); + Job updatedJob = jobService.failJob(job, "Async processing failed: " + e.getMessage()); + log.info("Job {} marked as FAILED from async wrapper, new status: {}", + job.getExternalId(), updatedJob != null ? updatedJob.getStatus() : "null"); + } catch (Exception failError) { + log.error("Failed to mark job {} as failed from async wrapper: {}", + job.getExternalId(), failError.getMessage(), failError); + } + } + } + + /** + * Load project and eagerly initialize all lazy associations within a short read-only transaction. + * This is separated from the main processing so the DB transaction is NOT held open + * during the potentially long-running AI analysis (which can take 10+ minutes for large PRs). + */ + @Transactional(timeout = 30, readOnly = true) + public Project loadAndInitializeProject(Long projectId) { + Project project = projectRepository.findById(projectId) + .orElseThrow(() -> new IllegalStateException("Project not found: " + projectId)); + initializeProjectAssociations(project); + return project; + } + + /** + * Process webhook - called from async method via self-injection. + * + * IMPORTANT: This method is intentionally NOT @Transactional. + * All downstream services (JobService, PullRequestService, CodeAnalysisService, etc.) + * manage their own transaction boundaries. Wrapping the entire webhook processing in a + * single transaction caused timeout failures on large PRs where AI analysis takes 10+ minutes, + * resulting in lost results after tokens were already spent. + * + * Project loading uses a short read-only transaction via {@link #loadAndInitializeProject(Long)}. + */ + public void processWebhookInTransaction( + EVcsProvider provider, + Long projectId, + WebhookPayload payload, + WebhookHandler handler, + Job job + ) { + log.info("processWebhookInTransaction ENTERED for job {}", job.getExternalId()); String placeholderCommentId = null; Project project = null; try { - // Re-fetch project within transaction to ensure all lazy associations are available - project = projectRepository.findById(projectId) - .orElseThrow(() -> new IllegalStateException("Project not found: " + projectId)); - - // Initialize lazy associations we'll need - initializeProjectAssociations(project); + // Load project in a short read-only transaction + // All lazy associations are eagerly initialized before the transaction closes + project = self.loadAndInitializeProject(projectId); + log.info("Calling jobService.startJob for job {}", job.getExternalId()); jobService.startJob(job); + log.info("jobService.startJob completed for job {}", job.getExternalId()); // Post placeholder comment immediately if this is a CodeCrow command on a PR if (payload.hasCodecrowCommand() && payload.pullRequestId() != null) { @@ -113,21 +140,25 @@ public void processWebhookAsync( final String finalPlaceholderCommentId = placeholderCommentId; // Create event consumer that logs to job + log.info("Calling handler.handle for job {}", job.getExternalId()); WebhookHandler.WebhookResult result = handler.handle(payload, project, event -> { String message = (String) event.getOrDefault("message", "Processing..."); String state = (String) event.getOrDefault("state", "processing"); jobService.info(job, state, message); }); + log.info("handler.handle completed for job {}, result status={}", job.getExternalId(), result.status()); // Check if the webhook was ignored (e.g., branch not matching pattern, analysis disabled) if ("ignored".equals(result.status())) { - log.info("Webhook ignored: {}", result.message()); + log.info("Webhook ignored for job {}: {}", job.getExternalId(), result.message()); // Delete placeholder if we posted one for an ignored command if (finalPlaceholderCommentId != null) { deletePlaceholderComment(provider, project, payload, finalPlaceholderCommentId); } - // Delete the job entirely - don't clutter DB with ignored webhooks - jobService.deleteIgnoredJob(job, result.message()); + // Mark job as SKIPPED - simpler and more reliable than deletion + // which can have transaction/lock issues with concurrent requests + jobService.skipJob(job, result.message()); + log.info("Marked ignored job {} as SKIPPED", job.getExternalId()); return; } @@ -139,7 +170,9 @@ public void processWebhookAsync( Long analysisId = ((Number) result.data().get("analysisId")).longValue(); jobService.info(job, "complete", "Analysis completed. Analysis ID: " + analysisId); } + log.info("Calling jobService.completeJob for job {}", job.getExternalId()); jobService.completeJob(job); + log.info("jobService.completeJob completed for job {}", job.getExternalId()); } else { // Post error to VCS (update placeholder if exists) - but ensure failJob is always called try { @@ -151,31 +184,102 @@ public void processWebhookAsync( jobService.failJob(job, result.message()); } + } catch (DiffTooLargeException diffEx) { + // Handle diff too large - this is a soft skip, not an error + log.warn("Diff too large for analysis - skipping: {}", diffEx.getMessage()); + + String skipMessage = String.format( + "⚠️ **Analysis Skipped - PR Too Large**\n\n" + + "This PR's diff exceeds the configured token limit:\n" + + "- **Estimated tokens:** %,d\n" + + "- **Maximum allowed:** %,d (%.1f%% of limit)\n\n" + + "To analyze this PR, consider:\n" + + "1. Breaking it into smaller PRs\n" + + "2. Increasing the token limit in project settings\n" + + "3. Using `/codecrow analyze` command on specific commits", + diffEx.getEstimatedTokens(), + diffEx.getMaxAllowedTokens(), + diffEx.getUtilizationPercentage() + ); + + try { + if (project == null) { + project = self.loadAndInitializeProject(projectId); + } + postInfoToVcs(provider, project, payload, skipMessage, placeholderCommentId, job); + } catch (Exception postError) { + log.error("Failed to post skip message to VCS: {}", postError.getMessage()); + } + + try { + jobService.skipJob(job, "Diff too large: " + diffEx.getEstimatedTokens() + " tokens > " + diffEx.getMaxAllowedTokens() + " limit"); + } catch (Exception skipError) { + log.error("Failed to skip job: {}", skipError.getMessage()); + } + + } catch (AnalysisLockedException lockEx) { + // Handle lock acquisition failure - mark job as failed + log.warn("Lock acquisition failed for analysis: {}", lockEx.getMessage()); + + String failMessage = String.format( + "⚠️ **Analysis Failed - Resource Locked**\n\n" + + "Could not acquire analysis lock after timeout:\n" + + "- **Lock type:** %s\n" + + "- **Branch:** %s\n" + + "- **Project:** %d\n\n" + + "Another analysis may be in progress. Please try again later.", + lockEx.getLockType(), + lockEx.getBranchName(), + lockEx.getProjectId() + ); + + try { + if (project == null) { + project = self.loadAndInitializeProject(projectId); + } + postErrorToVcs(provider, project, payload, failMessage, placeholderCommentId, job); + } catch (Exception postError) { + log.error("Failed to post lock error to VCS: {}", postError.getMessage()); + } + + try { + log.info("Marking job {} as FAILED due to lock acquisition timeout", job.getExternalId()); + Job updatedJob = jobService.failJob(job, "Lock acquisition timeout: " + lockEx.getMessage()); + log.info("Job {} marked as FAILED, new status: {}", job.getExternalId(), updatedJob.getStatus()); + } catch (Exception failError) { + log.error("Failed to fail job {}: {}", job.getExternalId(), failError.getMessage(), failError); + } + // Return early after handling lock exception to avoid any further processing + return; + } catch (Exception e) { log.error("Error processing webhook for job {}", job.getExternalId(), e); try { if (project == null) { - project = projectRepository.findById(projectId).orElse(null); - } - if (project != null) { - initializeProjectAssociations(project); - postErrorToVcs(provider, project, payload, "Processing failed: " + e.getMessage(), placeholderCommentId, job); + project = self.loadAndInitializeProject(projectId); } + postErrorToVcs(provider, project, payload, "Processing failed: " + e.getMessage(), placeholderCommentId, job); } catch (Exception postError) { log.error("Failed to post error to VCS: {}", postError.getMessage()); } try { - jobService.failJob(job, "Processing failed: " + e.getMessage()); + log.info("Marking job {} as FAILED due to processing error: {}", job.getExternalId(), e.getMessage()); + Job updatedJob = jobService.failJob(job, "Processing failed: " + e.getMessage()); + log.info("Job {} marked as FAILED, new status: {}", job.getExternalId(), updatedJob.getStatus()); } catch (Exception failError) { - log.error("Failed to mark job as failed: {}", failError.getMessage()); + log.error("Failed to mark job {} as failed: {}", job.getExternalId(), failError.getMessage(), failError); } } } /** - * Initialize lazy associations that will be needed for VCS operations. + * Initialize lazy associations that will be needed during webhook processing. + * Must be called within an active Hibernate session (i.e., inside a @Transactional method). + * + * Touches all lazy proxies so they are fully loaded before the session closes, + * allowing downstream services to access them outside a transaction. */ private void initializeProjectAssociations(Project project) { // Force initialization of VCS connections using unified accessor @@ -188,6 +292,30 @@ private void initializeProjectAssociations(Project project) { vcsConn.getProviderType(); } } + + // Force initialization of Workspace (lazy @ManyToOne) — accessed by all AiClientServices + // when building analysis requests via project.getWorkspace().getName() + var workspace = project.getWorkspace(); + if (workspace != null) { + workspace.getName(); + } + + // Force initialization of AI binding chain (lazy @OneToOne) — accessed by all AiClientServices + // via project.getAiBinding().getAiConnection() to get provider config, model, API key + var aiBinding = project.getAiBinding(); + if (aiBinding != null) { + var aiConn = aiBinding.getAiConnection(); + if (aiConn != null) { + aiConn.getProviderKey(); + } + } + + // Force initialization of Quality Gate (lazy @ManyToOne) — accessed by + // CodeAnalysisService.getQualityGateForAnalysis() after session closes + var qualityGate = project.getQualityGate(); + if (qualityGate != null) { + qualityGate.getName(); + } } /** @@ -300,7 +428,7 @@ private void postAskReply(VcsReportingService reportingService, Project project, private void postWithMarker(VcsReportingService reportingService, Project project, WebhookPayload payload, String content, String commandType, String placeholderCommentId, Job job) throws IOException { - String marker = getMarkerForCommandType(commandType); + String marker = CommentPlaceholders.getMarkerForCommandType(commandType); // If we have a placeholder comment, update it instead of creating a new one if (placeholderCommentId != null) { @@ -371,7 +499,7 @@ private void postErrorToVcs(EVcsProvider provider, Project project, WebhookPaylo Long.parseLong(payload.pullRequestId()), placeholderCommentId, content, - CODECROW_COMMAND_MARKER + CommentPlaceholders.CODECROW_COMMAND_MARKER ); log.info("Updated placeholder comment {} with error for PR {}", placeholderCommentId, payload.pullRequestId()); } else { @@ -380,7 +508,7 @@ private void postErrorToVcs(EVcsProvider provider, Project project, WebhookPaylo project, Long.parseLong(payload.pullRequestId()), content, - CODECROW_COMMAND_MARKER + CommentPlaceholders.CODECROW_COMMAND_MARKER ); log.info("Posted error to PR {}", payload.pullRequestId()); } @@ -390,6 +518,45 @@ private void postErrorToVcs(EVcsProvider provider, Project project, WebhookPaylo } } + /** + * Post an info message to VCS as a comment (for skipped/info scenarios). + * If placeholderCommentId is provided, update that comment with the info. + */ + private void postInfoToVcs(EVcsProvider provider, Project project, WebhookPayload payload, + String infoMessage, String placeholderCommentId, Job job) { + try { + if (payload.pullRequestId() == null) { + return; + } + + VcsReportingService reportingService = vcsServiceFactory.getReportingService(provider); + + // If we have a placeholder comment, update it with the info + if (placeholderCommentId != null) { + reportingService.updateComment( + project, + Long.parseLong(payload.pullRequestId()), + placeholderCommentId, + infoMessage, + CommentPlaceholders.CODECROW_COMMAND_MARKER + ); + log.info("Updated placeholder comment {} with info message for PR {}", placeholderCommentId, payload.pullRequestId()); + } else { + // No placeholder - post new info comment + reportingService.postComment( + project, + Long.parseLong(payload.pullRequestId()), + infoMessage, + CommentPlaceholders.CODECROW_COMMAND_MARKER + ); + log.info("Posted info message to PR {}", payload.pullRequestId()); + } + + } catch (Exception e) { + log.error("Failed to post info to VCS: {}", e.getMessage()); + } + } + /** * Sanitize error messages for display on VCS platforms. * Removes sensitive technical details like API keys, quotas, and internal stack traces. @@ -483,8 +650,8 @@ private String postPlaceholderComment(EVcsProvider provider, Project project, ? payload.getCodecrowCommand().type().name().toLowerCase() : "default"; - String placeholderContent = getPlaceholderMessage(commandType); - String marker = getMarkerForCommandType(commandType); + String placeholderContent = CommentPlaceholders.getPlaceholderMessage(commandType); + String marker = CommentPlaceholders.getMarkerForCommandType(commandType); // Delete any previous comments with the same marker before posting placeholder try { @@ -533,28 +700,4 @@ private void deletePlaceholderComment(EVcsProvider provider, Project project, log.warn("Failed to delete placeholder comment {}: {}", commentId, e.getMessage()); } } - - /** - * Get the placeholder message for a command type. - */ - private String getPlaceholderMessage(String commandType) { - return switch (commandType.toLowerCase()) { - case "analyze" -> PLACEHOLDER_ANALYZE; - case "summarize" -> PLACEHOLDER_SUMMARIZE; - case "review" -> PLACEHOLDER_REVIEW; - case "ask" -> PLACEHOLDER_ASK; - default -> PLACEHOLDER_DEFAULT; - }; - } - - /** - * Get the comment marker for a command type. - */ - private String getMarkerForCommandType(String commandType) { - return switch (commandType.toLowerCase()) { - case "summarize" -> CODECROW_SUMMARY_MARKER; - case "review" -> CODECROW_REVIEW_MARKER; - default -> CODECROW_COMMAND_MARKER; - }; - } } diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/AskCommandProcessor.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/AskCommandProcessor.java index b8ceb92f..68c5f87d 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/AskCommandProcessor.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/AskCommandProcessor.java @@ -375,7 +375,7 @@ private AskRequest buildAskRequest( credentials.oAuthClient(), credentials.oAuthSecret(), credentials.accessToken(), - aiConnection.getTokenLimitation(), + project.getEffectiveConfig().maxAnalysisTokenLimit(), credentials.vcsProviderString(), analysisContext, context.issueReferences() diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/ReviewCommandProcessor.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/ReviewCommandProcessor.java index 78a8dcb2..92a00954 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/ReviewCommandProcessor.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/ReviewCommandProcessor.java @@ -154,6 +154,11 @@ private ReviewRequest buildReviewRequest(Project project, WebhookPayload payload } AIConnection aiConnection = project.getAiBinding().getAiConnection(); + + // CRITICAL: Log the AI connection being used for debugging + log.info("Building review command request for project={}, AI model={}, provider={}, aiConnectionId={}", + project.getId(), aiConnection.getAiModel(), aiConnection.getProviderKey(), aiConnection.getId()); + String decryptedApiKey = tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted()); // Get VCS credentials using centralized extractor @@ -180,7 +185,7 @@ private ReviewRequest buildReviewRequest(Project project, WebhookPayload payload credentials.oAuthClient(), credentials.oAuthSecret(), credentials.accessToken(), - aiConnection.getTokenLimitation(), + project.getEffectiveConfig().maxAnalysisTokenLimit(), credentials.vcsProviderString() ); diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/SummarizeCommandProcessor.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/SummarizeCommandProcessor.java index c2312413..8e3f992b 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/SummarizeCommandProcessor.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/processor/command/SummarizeCommandProcessor.java @@ -292,7 +292,7 @@ private SummarizeRequest buildSummarizeRequest( credentials.oAuthSecret(), credentials.accessToken(), diagramType == PrSummarizeCache.DiagramType.MERMAID, - aiConnection.getTokenLimitation(), + project.getEffectiveConfig().maxAnalysisTokenLimit(), credentials.vcsProviderString() ); diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/CommentCommandRateLimitService.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/CommentCommandRateLimitService.java index 3a5bfabe..fc26ca02 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/CommentCommandRateLimitService.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/CommentCommandRateLimitService.java @@ -54,6 +54,7 @@ public boolean isCommandAllowed(Project project) { /** * Record a command execution for rate limiting purposes. + * Uses atomic upsert to avoid race conditions with concurrent requests. * * @param project The project that executed the command */ @@ -68,17 +69,8 @@ public void recordCommand(Project project) { .truncatedTo(ChronoUnit.HOURS) .plus((OffsetDateTime.now().getMinute() / windowMinutes) * windowMinutes, ChronoUnit.MINUTES); - Optional existingRecord = rateLimitRepository - .findByProjectIdAndWindowStart(project.getId(), windowStart); - - if (existingRecord.isPresent()) { - existingRecord.get().incrementCommandCount(); - rateLimitRepository.save(existingRecord.get()); - } else { - CommentCommandRateLimit newRecord = new CommentCommandRateLimit(project, windowStart); - newRecord.incrementCommandCount(); - rateLimitRepository.save(newRecord); - } + // Use atomic upsert to avoid duplicate key violations from concurrent requests + rateLimitRepository.upsertCommandCount(project.getId(), windowStart); log.debug("Recorded command for project {}, window starting at {}", project.getId(), windowStart); } diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/WebhookDeduplicationService.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/WebhookDeduplicationService.java new file mode 100644 index 00000000..c24909c2 --- /dev/null +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/service/WebhookDeduplicationService.java @@ -0,0 +1,96 @@ +package org.rostilos.codecrow.pipelineagent.generic.service; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.stereotype.Service; + +import java.time.Instant; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Service to deduplicate webhook events based on commit hash. + * + * When a PR is merged in Bitbucket, it sends both: + * - pullrequest:fulfilled (with merge commit) + * - repo:push (with same merge commit) + * + * Both events would trigger analysis for the same commit, causing duplicate processing. + * This service tracks recently analyzed commits and skips duplicates within a time window. + */ +@Service +public class WebhookDeduplicationService { + + private static final Logger log = LoggerFactory.getLogger(WebhookDeduplicationService.class); + + /** + * Time window in seconds to consider events as duplicates. + */ + private static final long DEDUP_WINDOW_SECONDS = 30; + + /** + * Counter to throttle cleanup operations (run every N requests instead of every request). + */ + private final java.util.concurrent.atomic.AtomicInteger cleanupCounter = new java.util.concurrent.atomic.AtomicInteger(0); + private static final int CLEANUP_INTERVAL = 10; + + /** + * Cache of recently analyzed commits. + * Key: "projectId:commitHash" + * Value: timestamp when the analysis was triggered + */ + private final Map recentCommitAnalyses = new ConcurrentHashMap<>(); + + /** + * Check if a commit analysis should be skipped as a duplicate. + * If not a duplicate, records this commit for future deduplication. + * + * @param projectId The project ID + * @param commitHash The commit being analyzed + * @param eventType The webhook event type (for logging) + * @return true if this is a duplicate and should be skipped, false if it should proceed + */ + public boolean isDuplicateCommitAnalysis(Long projectId, String commitHash, String eventType) { + if (commitHash == null || commitHash.isBlank()) { + return false; + } + + String key = projectId + ":" + commitHash; + Instant now = Instant.now(); + + // Use putIfAbsent for atomic check-and-set to prevent race conditions + // where multiple threads could pass the duplicate check simultaneously + Instant existingTimestamp = recentCommitAnalyses.putIfAbsent(key, now); + + if (existingTimestamp != null) { + long secondsSinceLastAnalysis = now.getEpochSecond() - existingTimestamp.getEpochSecond(); + + if (secondsSinceLastAnalysis < DEDUP_WINDOW_SECONDS) { + log.info("Skipping duplicate commit analysis: project={}, commit={}, event={}, " + + "lastAnalysis={}s ago (within {}s window)", + projectId, commitHash, eventType, secondsSinceLastAnalysis, DEDUP_WINDOW_SECONDS); + return true; + } + + // Entry exists but is stale, update it + recentCommitAnalyses.put(key, now); + } + + // Cleanup old entries periodically (not on every request) + if (cleanupCounter.incrementAndGet() % CLEANUP_INTERVAL == 0) { + cleanupOldEntries(now); + } + + return false; + } + + /** + * Remove entries older than the dedup window to prevent memory growth. + */ + private void cleanupOldEntries(Instant now) { + recentCommitAnalyses.entrySet().removeIf(entry -> { + long age = now.getEpochSecond() - entry.getValue().getEpochSecond(); + return age > DEDUP_WINDOW_SECONDS * 2; + }); + } +} diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/utils/CommentPlaceholders.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/utils/CommentPlaceholders.java new file mode 100644 index 00000000..3c02b322 --- /dev/null +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/utils/CommentPlaceholders.java @@ -0,0 +1,72 @@ +package org.rostilos.codecrow.pipelineagent.generic.utils; + +public class CommentPlaceholders { + /** Comment markers for CodeCrow command responses */ + public static final String CODECROW_COMMAND_MARKER = ""; + public static final String CODECROW_SUMMARY_MARKER = ""; + public static final String CODECROW_REVIEW_MARKER = ""; + + /** Placeholder messages for commands */ + public static final String PLACEHOLDER_ANALYZE = """ + 🔄 **CodeCrow is analyzing this PR...** + + This may take a few minutes depending on the size of the changes. + I'll update this comment with the results when the analysis is complete. + """; + + public static final String PLACEHOLDER_SUMMARIZE = """ + 🔄 **CodeCrow is generating a summary...** + + I'm analyzing the changes and creating diagrams. + This comment will be updated with the summary when ready. + """; + + public static final String PLACEHOLDER_REVIEW = """ + 🔄 **CodeCrow is reviewing this PR...** + + I'm examining the code changes for potential issues. + This comment will be updated with the review results when complete. + """; + + public static final String PLACEHOLDER_ASK = """ + 🔄 **CodeCrow is processing your question...** + + I'm analyzing the context to provide a helpful answer. + """; + + public static final String PLACEHOLDER_DEFAULT = """ + 🔄 **CodeCrow is processing...** + + Please wait while I complete this task. + """; + + /** + * Get the placeholder message for a command type. + */ + public static String getPlaceholderMessage(String commandType) { + if (commandType == null) { + return PLACEHOLDER_DEFAULT; + } + return switch (commandType.toLowerCase()) { + case "analyze" -> PLACEHOLDER_ANALYZE; + case "summarize" -> PLACEHOLDER_SUMMARIZE; + case "review" -> PLACEHOLDER_REVIEW; + case "ask" -> PLACEHOLDER_ASK; + default -> PLACEHOLDER_DEFAULT; + }; + } + + /** + * Get the comment marker for a command type. + */ + public static String getMarkerForCommandType(String commandType) { + if (commandType == null) { + return CODECROW_COMMAND_MARKER; + } + return switch (commandType.toLowerCase()) { + case "summarize" -> CODECROW_SUMMARY_MARKER; + case "review" -> CODECROW_REVIEW_MARKER; + default -> CODECROW_COMMAND_MARKER; + }; + } +} diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/AbstractWebhookHandler.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/AbstractWebhookHandler.java index 989c4438..e4f5e0e3 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/AbstractWebhookHandler.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/AbstractWebhookHandler.java @@ -3,7 +3,6 @@ import org.rostilos.codecrow.core.model.codeanalysis.AnalysisType; import org.rostilos.codecrow.core.model.project.Project; import org.rostilos.codecrow.core.model.project.config.BranchAnalysisConfig; -import org.rostilos.codecrow.core.model.project.config.ProjectConfig; import org.rostilos.codecrow.core.util.BranchPatternMatcher; import java.util.List; diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/CommentCommandWebhookHandler.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/CommentCommandWebhookHandler.java index 6267cbd7..0a120fb7 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/CommentCommandWebhookHandler.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/generic/webhookhandler/CommentCommandWebhookHandler.java @@ -13,6 +13,8 @@ import org.rostilos.codecrow.core.persistence.repository.codeanalysis.PrSummarizeCacheRepository; import org.rostilos.codecrow.core.service.CodeAnalysisService; import org.rostilos.codecrow.analysisengine.dto.request.processor.PrProcessRequest; +import org.rostilos.codecrow.analysisengine.exception.AnalysisLockedException; +import org.rostilos.codecrow.analysisengine.exception.DiffTooLargeException; import org.rostilos.codecrow.analysisengine.processor.analysis.PullRequestAnalysisProcessor; import org.rostilos.codecrow.pipelineagent.generic.dto.webhook.WebhookPayload; import org.rostilos.codecrow.pipelineagent.generic.service.CommandAuthorizationService; @@ -445,6 +447,14 @@ private WebhookResult runPrAnalysis( // If we got here, the processor posted results directly (which it does) return WebhookResult.success("Analysis completed", Map.of("commandType", commandType)); + } catch (DiffTooLargeException e) { + // Re-throw DiffTooLargeException so WebhookAsyncProcessor can handle it with proper job status + log.warn("PR diff too large for {} command: {}", commandType, e.getMessage()); + throw e; + } catch (AnalysisLockedException e) { + // Re-throw AnalysisLockedException so WebhookAsyncProcessor can handle it with proper job status + log.warn("Lock acquisition failed for {} command: {}", commandType, e.getMessage()); + throw e; } catch (Exception e) { log.error("Error running PR analysis for {} command: {}", commandType, e.getMessage(), e); return WebhookResult.error("Analysis failed: " + e.getMessage()); diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/service/GitHubAiClientService.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/service/GitHubAiClientService.java index 7075aca6..fcd28019 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/service/GitHubAiClientService.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/service/GitHubAiClientService.java @@ -2,6 +2,8 @@ import com.fasterxml.jackson.databind.JsonNode; import okhttp3.OkHttpClient; +import org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment.PrEnrichmentDataDto; +import org.rostilos.codecrow.analysisengine.service.PrFileEnrichmentService; import org.rostilos.codecrow.core.model.ai.AIConnection; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisMode; import org.rostilos.codecrow.core.model.codeanalysis.CodeAnalysis; @@ -13,10 +15,13 @@ import org.rostilos.codecrow.analysisengine.dto.request.processor.AnalysisProcessRequest; import org.rostilos.codecrow.analysisengine.dto.request.processor.BranchProcessRequest; import org.rostilos.codecrow.analysisengine.dto.request.processor.PrProcessRequest; +import org.rostilos.codecrow.analysisengine.exception.DiffTooLargeException; import org.rostilos.codecrow.analysisengine.service.vcs.VcsAiClientService; import org.rostilos.codecrow.analysisengine.util.DiffContentFilter; import org.rostilos.codecrow.analysisengine.util.DiffParser; +import org.rostilos.codecrow.analysisengine.util.TokenEstimator; import org.rostilos.codecrow.security.oauth.TokenEncryptionService; +import org.rostilos.codecrow.vcsclient.VcsClient; import org.rostilos.codecrow.vcsclient.VcsClientProvider; import org.rostilos.codecrow.vcsclient.github.actions.GetCommitRangeDiffAction; import org.rostilos.codecrow.vcsclient.github.actions.GetPullRequestAction; @@ -25,6 +30,7 @@ import org.rostilos.codecrow.vcsclient.utils.VcsConnectionCredentialsExtractor.VcsConnectionCredentials; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Service; import java.io.IOException; @@ -51,14 +57,17 @@ public class GitHubAiClientService implements VcsAiClientService { private final TokenEncryptionService tokenEncryptionService; private final VcsClientProvider vcsClientProvider; private final VcsConnectionCredentialsExtractor credentialsExtractor; + private final PrFileEnrichmentService enrichmentService; public GitHubAiClientService( TokenEncryptionService tokenEncryptionService, - VcsClientProvider vcsClientProvider + VcsClientProvider vcsClientProvider, + @Autowired(required = false) PrFileEnrichmentService enrichmentService ) { this.tokenEncryptionService = tokenEncryptionService; this.vcsClientProvider = vcsClientProvider; this.credentialsExtractor = new VcsConnectionCredentialsExtractor(tokenEncryptionService); + this.enrichmentService = enrichmentService; } @Override @@ -88,18 +97,38 @@ public AiAnalysisRequest buildAiAnalysisRequest( case BRANCH_ANALYSIS: return buildBranchAnalysisRequest(project, (BranchProcessRequest) request, previousAnalysis); default: - return buildPrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis); + return buildPrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis, Collections.emptyList()); + } + } + + @Override + public AiAnalysisRequest buildAiAnalysisRequest( + Project project, + AnalysisProcessRequest request, + Optional previousAnalysis, + List allPrAnalyses + ) throws GeneralSecurityException { + switch (request.getAnalysisType()) { + case BRANCH_ANALYSIS: + return buildBranchAnalysisRequest(project, (BranchProcessRequest) request, previousAnalysis); + default: + return buildPrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis, allPrAnalyses); } } private AiAnalysisRequest buildPrAnalysisRequest( Project project, PrProcessRequest request, - Optional previousAnalysis + Optional previousAnalysis, + List allPrAnalyses ) throws GeneralSecurityException { VcsInfo vcsInfo = getVcsInfo(project); VcsConnection vcsConnection = vcsInfo.vcsConnection(); AIConnection aiConnection = project.getAiBinding().getAiConnection(); + + // CRITICAL: Log the AI connection being used for debugging + log.info("Building PR analysis request for project={}, AI model={}, provider={}, aiConnectionId={}", + project.getId(), aiConnection.getAiModel(), aiConnection.getProviderKey(), aiConnection.getId()); // Initialize variables List changedFiles = Collections.emptyList(); @@ -149,6 +178,23 @@ private AiAnalysisRequest buildPrAnalysisRequest( originalSize > 0 ? (100 - (filteredSize * 100 / originalSize)) : 0); } + // Check token limit before proceeding with analysis + int maxTokenLimit = project.getEffectiveConfig().maxAnalysisTokenLimit(); + TokenEstimator.TokenEstimationResult tokenEstimate = TokenEstimator.estimateAndCheck(rawDiff, maxTokenLimit); + log.info("Token estimation for PR diff: {}", tokenEstimate.toLogString()); + + if (tokenEstimate.exceedsLimit()) { + log.warn("PR diff exceeds token limit - skipping analysis. Project={}, PR={}, Tokens={}/{}", + project.getId(), request.getPullRequestId(), + tokenEstimate.estimatedTokens(), tokenEstimate.maxAllowedTokens()); + throw new DiffTooLargeException( + tokenEstimate.estimatedTokens(), + tokenEstimate.maxAllowedTokens(), + project.getId(), + request.getPullRequestId() + ); + } + // Determine analysis mode: INCREMENTAL if we have previous analysis with different commit boolean canUseIncremental = previousAnalysis.isPresent() && previousCommitHash != null @@ -199,15 +245,35 @@ private AiAnalysisRequest buildPrAnalysisRequest( log.warn("Failed to fetch/parse PR metadata/diff for RAG context: {}", e.getMessage()); } - var builder = AiAnalysisRequestImpl.builder() + // Enrich PR with full file contents and dependency graph + PrEnrichmentDataDto enrichmentData = PrEnrichmentDataDto.empty(); + if (enrichmentService != null && enrichmentService.isEnrichmentEnabled() && !changedFiles.isEmpty()) { + try { + VcsClient vcsClient = vcsClientProvider.getClient(vcsConnection); + enrichmentData = enrichmentService.enrichPrFiles( + vcsClient, + vcsInfo.owner(), + vcsInfo.repoSlug(), + request.getSourceBranchName(), + changedFiles + ); + log.info("PR enrichment completed: {} files enriched, {} relationships", + enrichmentData.stats().filesEnriched(), + enrichmentData.stats().relationshipsFound()); + } catch (Exception e) { + log.warn("Failed to enrich PR files (non-critical): {}", e.getMessage()); + } + } + + AiAnalysisRequestImpl.Builder builder = AiAnalysisRequestImpl.builder() .withProjectId(project.getId()) .withPullRequestId(request.getPullRequestId()) .withProjectAiConnection(aiConnection) .withProjectVcsConnectionBindingInfo(vcsInfo.owner(), vcsInfo.repoSlug()) .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted())) .withUseLocalMcp(true) - .withPreviousAnalysisData(previousAnalysis) - .withMaxAllowedTokens(aiConnection.getTokenLimitation()) + .withAllPrAnalysesData(allPrAnalyses) // Use full PR history instead of just previous version + .withMaxAllowedTokens(project.getEffectiveConfig().maxAnalysisTokenLimit()) .withAnalysisType(request.getAnalysisType()) .withPrTitle(prTitle) .withPrDescription(prDescription) @@ -221,7 +287,9 @@ private AiAnalysisRequest buildPrAnalysisRequest( .withAnalysisMode(analysisMode) .withDeltaDiff(deltaDiff) .withPreviousCommitHash(previousCommitHash) - .withCurrentCommitHash(currentCommitHash); + .withCurrentCommitHash(currentCommitHash) + // File enrichment data + .withEnrichmentData(enrichmentData); addVcsCredentials(builder, vcsConnection); @@ -275,7 +343,7 @@ private AiAnalysisRequest buildBranchAnalysisRequest( .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted())) .withUseLocalMcp(true) .withPreviousAnalysisData(previousAnalysis) - .withMaxAllowedTokens(aiConnection.getTokenLimitation()) + .withMaxAllowedTokens(project.getEffectiveConfig().maxAnalysisTokenLimit()) .withAnalysisType(request.getAnalysisType()) .withTargetBranchName(request.getTargetBranchName()) .withCurrentCommitHash(request.getCommitHash()) @@ -287,7 +355,7 @@ private AiAnalysisRequest buildBranchAnalysisRequest( return builder.build(); } - private void addVcsCredentials(AiAnalysisRequestImpl.Builder builder, VcsConnection connection) + private void addVcsCredentials(AiAnalysisRequestImpl.Builder builder, VcsConnection connection) throws GeneralSecurityException { VcsConnectionCredentials credentials = credentialsExtractor.extractCredentials(connection); if (VcsConnectionCredentialsExtractor.hasAccessToken(credentials)) { diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/webhookhandler/GitHubPullRequestWebhookHandler.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/webhookhandler/GitHubPullRequestWebhookHandler.java index b64a4f7d..88618799 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/webhookhandler/GitHubPullRequestWebhookHandler.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/github/webhookhandler/GitHubPullRequestWebhookHandler.java @@ -1,10 +1,14 @@ package org.rostilos.codecrow.pipelineagent.github.webhookhandler; +import org.rostilos.codecrow.core.model.analysis.AnalysisLockType; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisType; import org.rostilos.codecrow.core.model.project.Project; import org.rostilos.codecrow.core.model.vcs.EVcsProvider; import org.rostilos.codecrow.analysisengine.dto.request.processor.PrProcessRequest; +import org.rostilos.codecrow.analysisengine.exception.AnalysisLockedException; +import org.rostilos.codecrow.analysisengine.exception.DiffTooLargeException; import org.rostilos.codecrow.analysisengine.processor.analysis.PullRequestAnalysisProcessor; +import org.rostilos.codecrow.analysisengine.service.AnalysisLockService; import org.rostilos.codecrow.analysisengine.service.vcs.VcsReportingService; import org.rostilos.codecrow.analysisengine.service.vcs.VcsServiceFactory; import org.rostilos.codecrow.pipelineagent.generic.dto.webhook.WebhookPayload; @@ -15,6 +19,7 @@ import org.springframework.stereotype.Component; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; @@ -47,13 +52,16 @@ public class GitHubPullRequestWebhookHandler extends AbstractWebhookHandler impl private final PullRequestAnalysisProcessor pullRequestAnalysisProcessor; private final VcsServiceFactory vcsServiceFactory; + private final AnalysisLockService analysisLockService; public GitHubPullRequestWebhookHandler( PullRequestAnalysisProcessor pullRequestAnalysisProcessor, - VcsServiceFactory vcsServiceFactory + VcsServiceFactory vcsServiceFactory, + AnalysisLockService analysisLockService ) { this.pullRequestAnalysisProcessor = pullRequestAnalysisProcessor; this.vcsServiceFactory = vcsServiceFactory; + this.analysisLockService = analysisLockService; } @Override @@ -115,8 +123,26 @@ private WebhookResult handlePullRequestEvent( Consumer> eventConsumer ) { String placeholderCommentId = null; + String acquiredLockKey = null; try { + // Try to acquire lock atomically BEFORE posting placeholder + // This prevents race condition where multiple webhooks could post duplicate placeholders + String sourceBranch = payload.sourceBranch(); + Optional earlyLock = analysisLockService.acquireLock( + project, sourceBranch, AnalysisLockType.PR_ANALYSIS, + payload.commitHash(), Long.parseLong(payload.pullRequestId())); + + if (earlyLock.isEmpty()) { + log.info("PR analysis already in progress for project={}, branch={}, PR={} - skipping duplicate webhook", + project.getId(), sourceBranch, payload.pullRequestId()); + return WebhookResult.ignored("PR analysis already in progress for this branch"); + } + + acquiredLockKey = earlyLock.get(); + + // Lock acquired - placeholder posting is now protected from race conditions + // Post placeholder comment immediately to show analysis has started placeholderCommentId = postPlaceholderComment(project, Long.parseLong(payload.pullRequestId())); @@ -131,6 +157,8 @@ private WebhookResult handlePullRequestEvent( request.placeholderCommentId = placeholderCommentId; request.prAuthorId = payload.prAuthorId(); request.prAuthorUsername = payload.prAuthorUsername(); + // Pass the pre-acquired lock key to avoid double-locking in the processor + request.preAcquiredLockKey = acquiredLockKey; log.info("Processing PR analysis: project={}, PR={}, source={}, target={}, placeholderCommentId={}", project.getId(), request.pullRequestId, request.sourceBranchName, request.targetBranchName, placeholderCommentId); @@ -156,8 +184,20 @@ private WebhookResult handlePullRequestEvent( return WebhookResult.success("PR analysis completed", result); + } catch (DiffTooLargeException | AnalysisLockedException e) { + // Re-throw these exceptions so WebhookAsyncProcessor can handle them properly + // Release the lock since processor won't take ownership + if (acquiredLockKey != null) { + analysisLockService.releaseLock(acquiredLockKey); + } + log.warn("PR analysis failed with recoverable exception for project {}: {}", project.getId(), e.getMessage()); + throw e; } catch (Exception e) { log.error("PR analysis failed for project {}", project.getId(), e); + // Release the lock since processor won't take ownership + if (acquiredLockKey != null) { + analysisLockService.releaseLock(acquiredLockKey); + } // Try to update placeholder with error message if (placeholderCommentId != null) { try { diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/service/GitLabAiClientService.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/service/GitLabAiClientService.java index a6dd0fce..e6b48c10 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/service/GitLabAiClientService.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/service/GitLabAiClientService.java @@ -2,6 +2,8 @@ import com.fasterxml.jackson.databind.JsonNode; import okhttp3.OkHttpClient; +import org.rostilos.codecrow.analysisengine.dto.request.ai.enrichment.PrEnrichmentDataDto; +import org.rostilos.codecrow.analysisengine.service.PrFileEnrichmentService; import org.rostilos.codecrow.core.model.ai.AIConnection; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisMode; import org.rostilos.codecrow.core.model.codeanalysis.CodeAnalysis; @@ -13,10 +15,13 @@ import org.rostilos.codecrow.analysisengine.dto.request.processor.AnalysisProcessRequest; import org.rostilos.codecrow.analysisengine.dto.request.processor.BranchProcessRequest; import org.rostilos.codecrow.analysisengine.dto.request.processor.PrProcessRequest; +import org.rostilos.codecrow.analysisengine.exception.DiffTooLargeException; import org.rostilos.codecrow.analysisengine.service.vcs.VcsAiClientService; import org.rostilos.codecrow.analysisengine.util.DiffContentFilter; import org.rostilos.codecrow.analysisengine.util.DiffParser; +import org.rostilos.codecrow.analysisengine.util.TokenEstimator; import org.rostilos.codecrow.security.oauth.TokenEncryptionService; +import org.rostilos.codecrow.vcsclient.VcsClient; import org.rostilos.codecrow.vcsclient.VcsClientProvider; import org.rostilos.codecrow.vcsclient.gitlab.actions.GetCommitRangeDiffAction; import org.rostilos.codecrow.vcsclient.gitlab.actions.GetMergeRequestAction; @@ -25,6 +30,7 @@ import org.rostilos.codecrow.vcsclient.utils.VcsConnectionCredentialsExtractor.VcsConnectionCredentials; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Service; import java.io.IOException; @@ -51,14 +57,17 @@ public class GitLabAiClientService implements VcsAiClientService { private final TokenEncryptionService tokenEncryptionService; private final VcsClientProvider vcsClientProvider; private final VcsConnectionCredentialsExtractor credentialsExtractor; + private final PrFileEnrichmentService enrichmentService; public GitLabAiClientService( TokenEncryptionService tokenEncryptionService, - VcsClientProvider vcsClientProvider + VcsClientProvider vcsClientProvider, + @Autowired(required = false) PrFileEnrichmentService enrichmentService ) { this.tokenEncryptionService = tokenEncryptionService; this.vcsClientProvider = vcsClientProvider; this.credentialsExtractor = new VcsConnectionCredentialsExtractor(tokenEncryptionService); + this.enrichmentService = enrichmentService; } @Override @@ -88,18 +97,38 @@ public AiAnalysisRequest buildAiAnalysisRequest( case BRANCH_ANALYSIS: return buildBranchAnalysisRequest(project, (BranchProcessRequest) request, previousAnalysis); default: - return buildMrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis); + return buildMrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis, Collections.emptyList()); + } + } + + @Override + public AiAnalysisRequest buildAiAnalysisRequest( + Project project, + AnalysisProcessRequest request, + Optional previousAnalysis, + List allPrAnalyses + ) throws GeneralSecurityException { + switch (request.getAnalysisType()) { + case BRANCH_ANALYSIS: + return buildBranchAnalysisRequest(project, (BranchProcessRequest) request, previousAnalysis); + default: + return buildMrAnalysisRequest(project, (PrProcessRequest) request, previousAnalysis, allPrAnalyses); } } private AiAnalysisRequest buildMrAnalysisRequest( Project project, PrProcessRequest request, - Optional previousAnalysis + Optional previousAnalysis, + List allPrAnalyses ) throws GeneralSecurityException { VcsInfo vcsInfo = getVcsInfo(project); VcsConnection vcsConnection = vcsInfo.vcsConnection(); AIConnection aiConnection = project.getAiBinding().getAiConnection(); + + // CRITICAL: Log the AI connection being used for debugging + log.info("Building MR analysis request for project={}, AI model={}, provider={}, aiConnectionId={}", + project.getId(), aiConnection.getAiModel(), aiConnection.getProviderKey(), aiConnection.getId()); // Initialize variables List changedFiles = Collections.emptyList(); @@ -150,6 +179,23 @@ private AiAnalysisRequest buildMrAnalysisRequest( originalSize > 0 ? (100 - (filteredSize * 100 / originalSize)) : 0); } + // Check token limit before proceeding with analysis + int maxTokenLimit = project.getEffectiveConfig().maxAnalysisTokenLimit(); + TokenEstimator.TokenEstimationResult tokenEstimate = TokenEstimator.estimateAndCheck(rawDiff, maxTokenLimit); + log.info("Token estimation for MR diff: {}", tokenEstimate.toLogString()); + + if (tokenEstimate.exceedsLimit()) { + log.warn("MR diff exceeds token limit - skipping analysis. Project={}, PR={}, Tokens={}/{}", + project.getId(), request.getPullRequestId(), + tokenEstimate.estimatedTokens(), tokenEstimate.maxAllowedTokens()); + throw new DiffTooLargeException( + tokenEstimate.estimatedTokens(), + tokenEstimate.maxAllowedTokens(), + project.getId(), + request.getPullRequestId() + ); + } + // Determine analysis mode: INCREMENTAL if we have previous analysis with different commit boolean canUseIncremental = previousAnalysis.isPresent() && previousCommitHash != null @@ -200,15 +246,35 @@ private AiAnalysisRequest buildMrAnalysisRequest( log.warn("Failed to fetch/parse MR metadata/diff for RAG context: {}", e.getMessage()); } - var builder = AiAnalysisRequestImpl.builder() + // Enrich PR with full file contents and dependency graph + PrEnrichmentDataDto enrichmentData = PrEnrichmentDataDto.empty(); + if (enrichmentService != null && enrichmentService.isEnrichmentEnabled() && !changedFiles.isEmpty()) { + try { + VcsClient vcsClient = vcsClientProvider.getClient(vcsConnection); + enrichmentData = enrichmentService.enrichPrFiles( + vcsClient, + vcsInfo.namespace(), + vcsInfo.repoSlug(), + request.getSourceBranchName(), + changedFiles + ); + log.info("PR enrichment completed: {} files enriched, {} relationships", + enrichmentData.stats().filesEnriched(), + enrichmentData.stats().relationshipsFound()); + } catch (Exception e) { + log.warn("Failed to enrich MR files (non-critical): {}", e.getMessage()); + } + } + + AiAnalysisRequestImpl.Builder builder = AiAnalysisRequestImpl.builder() .withProjectId(project.getId()) .withPullRequestId(request.getPullRequestId()) .withProjectAiConnection(aiConnection) .withProjectVcsConnectionBindingInfo(vcsInfo.namespace(), vcsInfo.repoSlug()) .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted())) .withUseLocalMcp(true) - .withPreviousAnalysisData(previousAnalysis) - .withMaxAllowedTokens(aiConnection.getTokenLimitation()) + .withAllPrAnalysesData(allPrAnalyses) // Use full PR history instead of just previous version + .withMaxAllowedTokens(project.getEffectiveConfig().maxAnalysisTokenLimit()) .withAnalysisType(request.getAnalysisType()) .withPrTitle(mrTitle) .withPrDescription(mrDescription) @@ -222,7 +288,9 @@ private AiAnalysisRequest buildMrAnalysisRequest( .withAnalysisMode(analysisMode) .withDeltaDiff(deltaDiff) .withPreviousCommitHash(previousCommitHash) - .withCurrentCommitHash(currentCommitHash); + .withCurrentCommitHash(currentCommitHash) + // File enrichment data + .withEnrichmentData(enrichmentData); addVcsCredentials(builder, vcsConnection); @@ -276,7 +344,7 @@ private AiAnalysisRequest buildBranchAnalysisRequest( .withProjectAiConnectionTokenDecrypted(tokenEncryptionService.decrypt(aiConnection.getApiKeyEncrypted())) .withUseLocalMcp(true) .withPreviousAnalysisData(previousAnalysis) - .withMaxAllowedTokens(aiConnection.getTokenLimitation()) + .withMaxAllowedTokens(project.getEffectiveConfig().maxAnalysisTokenLimit()) .withAnalysisType(request.getAnalysisType()) .withTargetBranchName(request.getTargetBranchName()) .withCurrentCommitHash(request.getCommitHash()) diff --git a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/webhookhandler/GitLabMergeRequestWebhookHandler.java b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/webhookhandler/GitLabMergeRequestWebhookHandler.java index 15fefd59..d3140156 100644 --- a/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/webhookhandler/GitLabMergeRequestWebhookHandler.java +++ b/java-ecosystem/services/pipeline-agent/src/main/java/org/rostilos/codecrow/pipelineagent/gitlab/webhookhandler/GitLabMergeRequestWebhookHandler.java @@ -1,10 +1,12 @@ package org.rostilos.codecrow.pipelineagent.gitlab.webhookhandler; +import org.rostilos.codecrow.core.model.analysis.AnalysisLockType; import org.rostilos.codecrow.core.model.codeanalysis.AnalysisType; import org.rostilos.codecrow.core.model.project.Project; import org.rostilos.codecrow.core.model.vcs.EVcsProvider; import org.rostilos.codecrow.analysisengine.dto.request.processor.PrProcessRequest; import org.rostilos.codecrow.analysisengine.processor.analysis.PullRequestAnalysisProcessor; +import org.rostilos.codecrow.analysisengine.service.AnalysisLockService; import org.rostilos.codecrow.analysisengine.service.vcs.VcsReportingService; import org.rostilos.codecrow.analysisengine.service.vcs.VcsServiceFactory; import org.rostilos.codecrow.pipelineagent.generic.dto.webhook.WebhookPayload; @@ -15,6 +17,7 @@ import org.springframework.stereotype.Component; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; @@ -47,13 +50,16 @@ public class GitLabMergeRequestWebhookHandler extends AbstractWebhookHandler imp private final PullRequestAnalysisProcessor pullRequestAnalysisProcessor; private final VcsServiceFactory vcsServiceFactory; + private final AnalysisLockService analysisLockService; public GitLabMergeRequestWebhookHandler( PullRequestAnalysisProcessor pullRequestAnalysisProcessor, - VcsServiceFactory vcsServiceFactory + VcsServiceFactory vcsServiceFactory, + AnalysisLockService analysisLockService ) { this.pullRequestAnalysisProcessor = pullRequestAnalysisProcessor; this.vcsServiceFactory = vcsServiceFactory; + this.analysisLockService = analysisLockService; } @Override @@ -117,8 +123,30 @@ private WebhookResult handleMergeRequestEvent( Consumer> eventConsumer ) { String placeholderCommentId = null; + String acquiredLockKey = null; try { + // Try to acquire lock atomically BEFORE posting placeholder + // This prevents TOCTOU race where multiple webhooks could pass isLocked() check simultaneously + // Note: PullRequestAnalysisProcessor.process() uses acquireLockWithWait() which will + // reuse this lock since it's for the same project/branch/type + String sourceBranch = payload.sourceBranch(); + Optional earlyLock = analysisLockService.acquireLock( + project, sourceBranch, AnalysisLockType.PR_ANALYSIS, + payload.commitHash(), Long.parseLong(payload.pullRequestId())); + + if (earlyLock.isEmpty()) { + log.info("MR analysis already in progress for project={}, branch={}, MR={} - skipping duplicate webhook", + project.getId(), sourceBranch, payload.pullRequestId()); + return WebhookResult.ignored("MR analysis already in progress for this branch"); + } + + acquiredLockKey = earlyLock.get(); + + // Lock acquired - placeholder posting is now protected from race conditions + // Note: We don't release this lock here - PullRequestAnalysisProcessor will manage it + // since acquireLockWithWait() will detect the existing lock and use it + // Post placeholder comment immediately to show analysis has started placeholderCommentId = postPlaceholderComment(project, Long.parseLong(payload.pullRequestId())); @@ -133,6 +161,8 @@ private WebhookResult handleMergeRequestEvent( request.placeholderCommentId = placeholderCommentId; request.prAuthorId = payload.prAuthorId(); request.prAuthorUsername = payload.prAuthorUsername(); + // Pass the pre-acquired lock key to avoid double-locking in the processor + request.preAcquiredLockKey = acquiredLockKey; log.info("Processing MR analysis: project={}, MR={}, source={}, target={}, placeholderCommentId={}", project.getId(), request.pullRequestId, request.sourceBranchName, request.targetBranchName, placeholderCommentId); @@ -159,6 +189,10 @@ private WebhookResult handleMergeRequestEvent( } catch (Exception e) { log.error("MR analysis failed for project {}", project.getId(), e); + // Release the lock since processor won't take ownership + if (acquiredLockKey != null) { + analysisLockService.releaseLock(acquiredLockKey); + } // Try to update placeholder with error message if (placeholderCommentId != null) { try { diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/controller/AIConnectionController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/controller/AIConnectionController.java index 1274b515..3bb2d788 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/controller/AIConnectionController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/controller/AIConnectionController.java @@ -5,6 +5,8 @@ import jakarta.validation.Valid; import org.rostilos.codecrow.core.model.ai.AIConnection; import org.rostilos.codecrow.core.model.workspace.Workspace; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.webserver.ai.dto.request.CreateAIConnectionRequest; import org.rostilos.codecrow.webserver.ai.dto.request.UpdateAiConnectionRequest; import org.rostilos.codecrow.core.dto.ai.AIConnectionDTO; @@ -12,7 +14,6 @@ import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import java.security.GeneralSecurityException; @@ -35,7 +36,7 @@ public AIConnectionController( } @GetMapping("/list") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity> listWorkspaceConnections(@PathVariable String workspaceSlug) { Workspace workspace = workspaceService.getWorkspaceBySlug(workspaceSlug); List workspaceConnections = aiConnectionService.listWorkspaceConnections(workspace.getId()) @@ -47,7 +48,7 @@ public ResponseEntity> listWorkspaceConnections(@PathVaria } @PostMapping("/create") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity createConnection( @PathVariable String workspaceSlug, @Valid @RequestBody CreateAIConnectionRequest request @@ -58,7 +59,7 @@ public ResponseEntity createConnection( } @PatchMapping("/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId, @@ -71,7 +72,7 @@ public ResponseEntity updateConnection( @DeleteMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity deleteConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/CreateAIConnectionRequest.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/CreateAIConnectionRequest.java index 558f236e..307dc44e 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/CreateAIConnectionRequest.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/CreateAIConnectionRequest.java @@ -13,6 +13,4 @@ public class CreateAIConnectionRequest { public String aiModel; @NotBlank(message = "API key is required") public String apiKey; - @NotBlank(message = "Please specify max token limit") - public String tokenLimitation; } diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/UpdateAiConnectionRequest.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/UpdateAiConnectionRequest.java index ae834bb5..f7805b90 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/UpdateAiConnectionRequest.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/dto/request/UpdateAiConnectionRequest.java @@ -9,5 +9,4 @@ public class UpdateAiConnectionRequest { public AIProviderKey providerKey; public String aiModel; public String apiKey; - public String tokenLimitation; } diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/service/AIConnectionService.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/service/AIConnectionService.java index 510f0901..48949fc7 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/service/AIConnectionService.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/ai/service/AIConnectionService.java @@ -53,7 +53,6 @@ public AIConnection createAiConnection(Long workspaceId, CreateAIConnectionReque newAiConnection.setProviderKey(request.providerKey); newAiConnection.setAiModel(request.aiModel); newAiConnection.setApiKeyEncrypted(apiKeyEncrypted); - newAiConnection.setTokenLimitation(Integer.parseInt(request.tokenLimitation)); return connectionRepository.save(newAiConnection); } @@ -77,9 +76,6 @@ public AIConnection updateAiConnection(Long workspaceId, Long connectionId, Upda connection.setApiKeyEncrypted(apiKeyEncrypted); } - if(request.tokenLimitation != null && !request.tokenLimitation.isEmpty()) { - connection.setTokenLimitation(Integer.parseInt(request.tokenLimitation)); - } return connectionRepository.save(connection); } diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/AnalysisIssueController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/AnalysisIssueController.java index 78c5d7d1..606cbf7d 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/AnalysisIssueController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/AnalysisIssueController.java @@ -4,6 +4,8 @@ import org.rostilos.codecrow.core.model.codeanalysis.CodeAnalysisIssue; import org.rostilos.codecrow.core.model.workspace.Workspace; import org.rostilos.codecrow.core.service.CodeAnalysisService; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.webserver.analysis.dto.request.IssueStatusUpdateRequest; import org.rostilos.codecrow.webserver.analysis.dto.response.IssueStatusUpdateResponse; import org.rostilos.codecrow.webserver.analysis.service.AnalysisService; @@ -16,7 +18,6 @@ import org.slf4j.LoggerFactory; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import java.util.List; @@ -52,7 +53,7 @@ public AnalysisIssueController( * Query params: ?branch={branch}&pullRequestId={prId}&severity={critical|high|medium|low|info}&type={security|quality|performance|style} */ @GetMapping - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity listIssues( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -101,7 +102,7 @@ public ResponseEntity listIssues( } @GetMapping("/{issueId}") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity getIssueById( @PathVariable Long issueId, @PathVariable String workspaceSlug, @@ -134,7 +135,7 @@ public ResponseEntity getIssueById( } @PutMapping("/{issueId}/status") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateIssueStatus( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/PullRequestController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/PullRequestController.java index 5af60ac6..663ca63d 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/PullRequestController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analysis/controller/PullRequestController.java @@ -11,6 +11,7 @@ import org.rostilos.codecrow.core.model.codeanalysis.CodeAnalysis; import org.rostilos.codecrow.core.persistence.repository.pullrequest.PullRequestRepository; import org.rostilos.codecrow.core.persistence.repository.codeanalysis.CodeAnalysisRepository; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.webserver.project.service.ProjectService; import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.rostilos.codecrow.core.model.project.Project; @@ -21,12 +22,12 @@ import org.rostilos.codecrow.core.dto.analysis.issue.IssueDTO; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import org.rostilos.codecrow.core.dto.pullrequest.PullRequestDTO; @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@IsWorkspaceMember @RequestMapping("/api/{workspaceSlug}/project/{projectNamespace}/pull-requests") public class PullRequestController { private final PullRequestRepository pullRequestRepository; @@ -49,7 +50,6 @@ public PullRequestController(PullRequestRepository pullRequestRepository, Projec } @GetMapping - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> listPullRequests( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -72,7 +72,6 @@ public ResponseEntity> listPullRequests( } @GetMapping("/by-branch") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity>> listPullRequestsByBranch( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -110,7 +109,6 @@ public ResponseEntity>> listPullRequestsByBranc } @GetMapping("/branches/issues") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> listBranchIssues( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analytics/controller/ProjectAnalyticsController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analytics/controller/ProjectAnalyticsController.java index ccbe0404..f6b87583 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analytics/controller/ProjectAnalyticsController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/analytics/controller/ProjectAnalyticsController.java @@ -4,6 +4,8 @@ import org.rostilos.codecrow.core.model.codeanalysis.CodeAnalysis; import org.rostilos.codecrow.core.model.codeanalysis.CodeAnalysisIssue; import org.rostilos.codecrow.core.model.codeanalysis.IssueSeverity; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.webserver.analysis.dto.response.AnalysesHistoryResponse; import org.rostilos.codecrow.webserver.analysis.dto.response.IssueStatusUpdateResponse; import org.rostilos.codecrow.webserver.analysis.service.AnalysisService; @@ -13,7 +15,6 @@ import org.rostilos.codecrow.core.service.CodeAnalysisService; import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import java.time.OffsetDateTime; @@ -21,6 +22,7 @@ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@IsWorkspaceMember @RequestMapping("/api/{workspaceSlug}/projects/{projectNamespace}/analysis") public class ProjectAnalyticsController { @@ -55,7 +57,6 @@ public ProjectAnalyticsController( * } */ @GetMapping("/summary") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getProjectSummary( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -97,7 +98,6 @@ public ResponseEntity getProjectSummary( * Response shape described in the task — this method builds that shape using available services. */ @GetMapping("/detailed-stats") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getDetailedStats( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -342,7 +342,6 @@ public ResponseEntity getDetailedStats( * Query params: ?page={page}&pageSize={size}&branch={branch} */ @GetMapping("/history") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getAnalysisHistory( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -368,7 +367,6 @@ public ResponseEntity getAnalysisHistory( * Returns a list of points {date, resolvedCount, totalIssues, resolvedRate} */ @GetMapping("/trends/resolved") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> getResolvedTrend( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -390,7 +388,6 @@ public ResponseEntity> getResol * Note: This endpoint requires a branch parameter and returns branch-specific data */ @GetMapping("/trends/issues") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> getBranchIssuesTrend( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -410,7 +407,7 @@ public ResponseEntity> getB * Update a single issue's status */ @PutMapping("/issues/{issueId}/status") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateIssueStatus( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -459,7 +456,7 @@ public ResponseEntity updateIssueStatus( * Update multiple issues' status at once */ @PutMapping("/issues/bulk-status") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity bulkUpdateIssueStatus( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/BitbucketConnectController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/BitbucketConnectController.java index 431748c7..a7e2ff26 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/BitbucketConnectController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/BitbucketConnectController.java @@ -4,6 +4,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import io.jsonwebtoken.Claims; import org.rostilos.codecrow.core.model.vcs.BitbucketConnectInstallation; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; import org.rostilos.codecrow.security.service.UserDetailsImpl; import org.rostilos.codecrow.webserver.integration.dto.response.VcsConnectionDTO; import org.rostilos.codecrow.webserver.exception.IntegrationException; @@ -65,7 +66,7 @@ public BitbucketConnectController(BitbucketConnectService connectService) { * The frontend will redirect the user to this URL. */ @PostMapping("/install/start") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> startInstall( @RequestParam Long workspaceId, @RequestParam(required = false) String workspaceSlug) { @@ -99,7 +100,7 @@ public ResponseEntity> startInstall( * Frontend polls this after redirecting user to Bitbucket. */ @GetMapping("/install/status") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> checkInstallStatus(@RequestParam String state) { PendingInstall pending = pendingInstalls.get(state); diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/VcsIntegrationController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/VcsIntegrationController.java index e890467b..d6c88a37 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/VcsIntegrationController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/integration/controller/VcsIntegrationController.java @@ -1,6 +1,7 @@ package org.rostilos.codecrow.webserver.integration.controller; import jakarta.validation.Valid; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; import org.rostilos.codecrow.webserver.generic.dto.message.MessageResponse; import org.rostilos.codecrow.core.model.vcs.EVcsConnectionType; import org.rostilos.codecrow.core.model.vcs.EVcsProvider; @@ -16,7 +17,6 @@ import org.slf4j.LoggerFactory; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import java.io.IOException; @@ -64,7 +64,7 @@ public ResponseEntity> handleIntegrationException(Integratio * GET /api/{workspaceSlug}/integrations/{provider}/app/install-url */ @GetMapping("/app/install-url") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity getInstallUrl( @PathVariable String workspaceSlug, @PathVariable String provider @@ -133,7 +133,7 @@ public ResponseEntity handleAppCallback( * @param connectionType Optional filter: APP, OAUTH_MANUAL, PERSONAL_TOKEN, etc. */ @GetMapping("/connections") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> getConnections( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -161,7 +161,7 @@ public ResponseEntity> getConnections( * GET /api/{workspaceSlug}/integrations/{provider}/connections/{connectionId} */ @GetMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity getConnection( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -180,7 +180,7 @@ public ResponseEntity getConnection( * DELETE /api/{workspaceSlug}/integrations/{provider}/connections/{connectionId} */ @DeleteMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity deleteConnection( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -199,7 +199,7 @@ public ResponseEntity deleteConnection( * POST /api/{workspaceSlug}/integrations/{provider}/connections/{connectionId}/sync */ @PostMapping("/connections/{connectionId}/sync") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity syncConnection( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -219,7 +219,7 @@ public ResponseEntity syncConnection( * GET /api/{workspaceSlug}/integrations/{provider}/connections/{connectionId}/reconnect-url */ @GetMapping("/connections/{connectionId}/reconnect-url") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity getReconnectUrl( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -238,7 +238,7 @@ public ResponseEntity getReconnectUrl( * GET /api/{workspaceSlug}/integrations/{provider}/repos */ @GetMapping("/repos") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity listRepositories( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -269,7 +269,7 @@ public ResponseEntity listRepositories( * @param limit Optional maximum number of results (default: 100, 0 for unlimited) */ @GetMapping("/repos/{externalRepoId}/branches") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> listBranches( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -299,7 +299,7 @@ public ResponseEntity> listBranches( * GET /api/{workspaceSlug}/integrations/{provider}/repos/{externalRepoId} */ @GetMapping("/repos/{externalRepoId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity getRepository( @PathVariable String workspaceSlug, @PathVariable String provider, @@ -326,7 +326,7 @@ public ResponseEntity getRepository( * POST /api/{workspaceSlug}/integrations/{provider}/repos/{externalRepoId}/onboard */ @PostMapping("/repos/{externalRepoId}/onboard") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity onboardRepository( @PathVariable String workspaceSlug, @PathVariable String provider, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/job/controller/JobController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/job/controller/JobController.java index 65bbdc9a..53ab078d 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/job/controller/JobController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/job/controller/JobController.java @@ -5,6 +5,7 @@ import org.rostilos.codecrow.core.model.project.Project; import org.rostilos.codecrow.core.persistence.repository.job.JobLogRepository; import org.rostilos.codecrow.core.service.JobService; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.webserver.project.service.ProjectService; import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.slf4j.Logger; @@ -14,7 +15,6 @@ import org.springframework.data.domain.Pageable; import org.springframework.http.MediaType; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; @@ -31,6 +31,7 @@ */ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@IsWorkspaceMember @RequestMapping("/api") public class JobController { @@ -59,7 +60,6 @@ public JobController( * List all jobs for a workspace. */ @GetMapping("/{workspaceSlug}/jobs") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity listWorkspaceJobs( @PathVariable String workspaceSlug, @RequestParam(defaultValue = "0") int page, @@ -89,7 +89,6 @@ public ResponseEntity listWorkspaceJobs( * List all jobs for a project. */ @GetMapping("/{workspaceSlug}/projects/{projectNamespace}/jobs") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity listProjectJobs( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -129,7 +128,6 @@ public ResponseEntity listProjectJobs( * Get active (running) jobs for a project. */ @GetMapping("/{workspaceSlug}/projects/{projectNamespace}/jobs/active") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> getActiveJobs( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -150,7 +148,6 @@ public ResponseEntity> getActiveJobs( * Get job details by external ID. */ @GetMapping("/{workspaceSlug}/projects/{projectNamespace}/jobs/{jobId}") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getJob( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -173,7 +170,6 @@ public ResponseEntity getJob( * Get job logs (paginated or all). */ @GetMapping("/{workspaceSlug}/projects/{projectNamespace}/jobs/{jobId}/logs") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getJobLogs( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -217,7 +213,6 @@ public ResponseEntity getJobLogs( */ @GetMapping(value = "/{workspaceSlug}/projects/{projectNamespace}/jobs/{jobId}/logs/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE) - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public SseEmitter streamJobLogs( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -304,7 +299,6 @@ public SseEmitter streamJobLogs( * Cancel a running job. */ @PostMapping("/{workspaceSlug}/projects/{projectNamespace}/jobs/{jobId}/cancel") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity cancelJob( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/AllowedCommandUserController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/AllowedCommandUserController.java index dd4e597f..5f9e0566 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/AllowedCommandUserController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/AllowedCommandUserController.java @@ -5,6 +5,7 @@ import org.rostilos.codecrow.core.model.project.AllowedCommandUser; import org.rostilos.codecrow.core.model.project.Project; import org.rostilos.codecrow.core.model.workspace.Workspace; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; import org.rostilos.codecrow.webserver.project.service.AllowedCommandUserService; import org.rostilos.codecrow.webserver.project.service.ProjectService; import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; @@ -12,7 +13,6 @@ import org.slf4j.LoggerFactory; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import java.time.OffsetDateTime; @@ -30,6 +30,7 @@ */ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@HasOwnerOrAdminRights @RequestMapping("/api/{workspaceSlug}/project/{projectNamespace}/allowed-users") public class AllowedCommandUserController { @@ -54,7 +55,6 @@ public AllowedCommandUserController( * Get all allowed command users for a project. */ @GetMapping - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getAllowedUsers( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -81,7 +81,6 @@ public ResponseEntity getAllowedUsers( * Add a user to the allowed list. */ @PostMapping - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity addAllowedUser( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -110,7 +109,6 @@ public ResponseEntity addAllowedUser( * Remove a user from the allowed list. */ @DeleteMapping("/{vcsUserId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity removeAllowedUser( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -130,7 +128,6 @@ public ResponseEntity removeAllowedUser( * Enable or disable a user. */ @PatchMapping("/{vcsUserId}/enabled") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity setUserEnabled( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -156,7 +153,6 @@ public ResponseEntity setUserEnabled( * Fetches repository collaborators with write access and adds them to the allowed list. */ @PostMapping("/sync") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity syncFromVcs( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -182,7 +178,6 @@ public ResponseEntity syncFromVcs( * Clear all allowed users for a project. */ @DeleteMapping - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity clearAllowedUsers( @PathVariable String workspaceSlug, @PathVariable String projectNamespace diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/ProjectController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/ProjectController.java index 9704f48b..77064a5d 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/ProjectController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/controller/ProjectController.java @@ -8,8 +8,9 @@ import org.rostilos.codecrow.core.model.project.Project; import org.rostilos.codecrow.core.model.project.config.CommentCommandsConfig; import org.rostilos.codecrow.core.model.project.config.InstallationMethod; -import org.rostilos.codecrow.core.model.project.config.ProjectConfig; import org.rostilos.codecrow.core.model.workspace.Workspace; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.security.service.UserDetailsImpl; import org.rostilos.codecrow.webserver.project.dto.request.BindAiConnectionRequest; import org.rostilos.codecrow.webserver.project.dto.request.BindRepositoryRequest; @@ -31,7 +32,6 @@ import org.springframework.http.HttpStatus; import org.springframework.http.MediaType; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.security.core.annotation.AuthenticationPrincipal; import org.springframework.web.bind.annotation.CrossOrigin; import org.springframework.web.bind.annotation.DeleteMapping; @@ -51,6 +51,7 @@ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@IsWorkspaceMember @RequestMapping("/api/{workspaceSlug}/project") public class ProjectController { private final ProjectService projectService; @@ -77,7 +78,6 @@ public ProjectController( } @GetMapping("/project_list") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> getUserWorkspaceProjectsList(@PathVariable String workspaceSlug) { Workspace workspace = workspaceService.getWorkspaceBySlug(workspaceSlug); List userWorkspaceProjects = projectService.listWorkspaceProjects(workspace.getId()); @@ -90,7 +90,6 @@ public ResponseEntity> getUserWorkspaceProjectsList(@PathVariab } @GetMapping("/projects") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> getProjectsPaginated( @PathVariable String workspaceSlug, @RequestParam(required = false, defaultValue = "") String search, @@ -125,7 +124,7 @@ public ResponseEntity> getProjectsPaginated( } @PostMapping("/create") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity createProject( @PathVariable String workspaceSlug, @Valid @RequestBody CreateProjectRequest request @@ -136,7 +135,7 @@ public ResponseEntity createProject( } @PostMapping("/{projectNamespace}/token/generate") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> generateProjectJwt( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -156,7 +155,7 @@ public ResponseEntity> generateProjectJwt( } @GetMapping("/{projectNamespace}/token") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> listProjectTokens( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -170,7 +169,7 @@ public ResponseEntity> listProjectTokens( } @DeleteMapping("/{projectNamespace}/token/{tokenId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity deleteProjectToken( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -185,7 +184,7 @@ public ResponseEntity deleteProjectToken( //TODO: service implementation, more fields to update @PatchMapping("/{projectNamespace}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateProject( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -198,7 +197,7 @@ public ResponseEntity updateProject( } @DeleteMapping("/{projectNamespace}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity deleteProject( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -214,7 +213,7 @@ public ResponseEntity deleteProject( } @PutMapping("/{projectNamespace}/repository/bind") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity bindRepository( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -227,7 +226,7 @@ public ResponseEntity bindRepository( } @DeleteMapping("/{projectNamespace}/repository/unbind") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity unbindRepository( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -244,7 +243,7 @@ public ResponseEntity unbindRepository( } @PatchMapping("/{projectNamespace}/repository/settings") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateRepositorySettings( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -257,7 +256,7 @@ public ResponseEntity updateRepositorySettings( } @PutMapping("/{projectNamespace}/ai/bind") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity bindAiConnection( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -274,7 +273,6 @@ public ResponseEntity bindAiConnection( * Returns list of analyzed branches for the project */ @GetMapping("/{projectNamespace}/branches") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity> getProjectBranches( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -303,7 +301,7 @@ public ResponseEntity> getProjectBranches( * Request body: { "branchId": 123 } or { "branchName": "main" } */ @PutMapping("/{projectNamespace}/default-branch") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity setDefaultBranch( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -344,7 +342,6 @@ public record SetDefaultBranchRequest( * Returns the branch analysis configuration for the project */ @GetMapping("/{projectNamespace}/branch-analysis-config") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getBranchAnalysisConfig( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -374,7 +371,7 @@ public record BranchAnalysisConfigResponse( * Request body: { "prTargetBranches": ["main", "develop", "release/*"], "branchPushPatterns": ["main", "develop"] } */ @PutMapping("/{projectNamespace}/branch-analysis-config") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateBranchAnalysisConfig( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -403,7 +400,6 @@ public record BranchAnalysisConfigRequest( * Returns the RAG indexing status for the project */ @GetMapping("/{projectNamespace}/rag/status") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getRagIndexStatus( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -432,7 +428,7 @@ public ResponseEntity getRagIndexStatus( * Updates the RAG configuration for the project (enable/disable, set branch, exclude patterns, delta config) */ @PutMapping("/{projectNamespace}/rag/config") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateRagConfig( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -462,7 +458,7 @@ public ResponseEntity updateRagConfig( * - Rate-limited to prevent abuse (min 60 seconds between triggers per project) */ @PostMapping(value = "/{projectNamespace}/rag/trigger", produces = MediaType.TEXT_EVENT_STREAM_VALUE) - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public SseEmitter triggerRagIndexing( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -524,7 +520,6 @@ public record RagStatusResponse( * Returns the comment commands configuration for the project */ @GetMapping("/{projectNamespace}/comment-commands-config") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") public ResponseEntity getCommentCommandsConfig( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -541,7 +536,7 @@ public ResponseEntity getCommentCommandsConfig( * Requires App-based VCS connection for comment commands to work. */ @PutMapping("/{projectNamespace}/comment-commands-config") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateCommentCommandsConfig( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -561,7 +556,7 @@ public ResponseEntity updateCommentCommandsConfig( * Updates analysis settings for the project (PR auto-analysis, branch analysis, installation method) */ @PutMapping("/{projectNamespace}/analysis-settings") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateAnalysisSettings( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -584,7 +579,8 @@ public ResponseEntity updateAnalysisSettings( project.getId(), request.prAnalysisEnabled(), request.branchAnalysisEnabled(), - installationMethod + installationMethod, + request.maxAnalysisTokenLimit() ); return new ResponseEntity<>(ProjectDTO.fromProject(updated), HttpStatus.OK); } @@ -592,14 +588,15 @@ public ResponseEntity updateAnalysisSettings( public record UpdateAnalysisSettingsRequest( Boolean prAnalysisEnabled, Boolean branchAnalysisEnabled, - String installationMethod + String installationMethod, + Integer maxAnalysisTokenLimit ) {} /** * Updates the quality gate for a project */ @PutMapping("/{projectNamespace}/quality-gate") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateProjectQualityGate( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, @@ -631,7 +628,7 @@ public record UpdateProjectQualityGateRequest( * - Webhook was accidentally deleted in the VCS provider */ @PostMapping("/{projectNamespace}/webhooks/setup") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity setupWebhooks( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -652,7 +649,7 @@ public ResponseEntity setupWebhooks( * Returns webhook configuration info for the project. */ @GetMapping("/{projectNamespace}/webhooks/info") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity getWebhookInfo( @PathVariable String workspaceSlug, @PathVariable String projectNamespace @@ -691,7 +688,7 @@ public record WebhookInfoResponse( * Warning: Changing VCS connection may require manual cleanup of old webhooks. */ @PutMapping("/{projectNamespace}/vcs-connection") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity changeVcsConnection( @PathVariable String workspaceSlug, @PathVariable String projectNamespace, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/IProjectService.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/IProjectService.java new file mode 100644 index 00000000..154b949e --- /dev/null +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/IProjectService.java @@ -0,0 +1,113 @@ +package org.rostilos.codecrow.webserver.project.service; + +import java.security.GeneralSecurityException; +import java.util.List; + +import org.rostilos.codecrow.core.model.branch.Branch; +import org.rostilos.codecrow.core.model.project.Project; +import org.rostilos.codecrow.core.model.project.config.BranchAnalysisConfig; +import org.rostilos.codecrow.core.model.project.config.CommentCommandsConfig; +import org.rostilos.codecrow.core.model.project.config.InstallationMethod; +import org.rostilos.codecrow.core.model.vcs.EVcsProvider; +import org.rostilos.codecrow.webserver.project.dto.request.BindAiConnectionRequest; +import org.rostilos.codecrow.webserver.project.dto.request.BindRepositoryRequest; +import org.rostilos.codecrow.webserver.project.dto.request.ChangeVcsConnectionRequest; +import org.rostilos.codecrow.webserver.project.dto.request.CreateProjectRequest; +import org.rostilos.codecrow.webserver.project.dto.request.UpdateCommentCommandsConfigRequest; +import org.rostilos.codecrow.webserver.project.dto.request.UpdateProjectRequest; +import org.rostilos.codecrow.webserver.project.dto.request.UpdateRepositorySettingsRequest; +import org.springframework.data.domain.Page; + +/** + * Interface for project management operations. + *

+ * This interface enables cloud implementations to extend or decorate + * the base project service with additional capabilities like billing limits. + */ +public interface IProjectService { + + // ==================== Core CRUD ==================== + + List listWorkspaceProjects(Long workspaceId); + + Page listWorkspaceProjectsPaginated(Long workspaceId, String search, int page, int size); + + Project createProject(Long workspaceId, CreateProjectRequest request); + + Project getProjectById(Long projectId); + + Project getProjectByWorkspaceAndNamespace(Long workspaceId, String namespace); + + Project updateProject(Long workspaceId, Long projectId, UpdateProjectRequest request); + + void deleteProject(Long workspaceId, Long projectId); + + void deleteProjectByNamespace(Long workspaceId, String namespace); + + // ==================== Repository Binding ==================== + + Project bindRepository(Long workspaceId, Long projectId, BindRepositoryRequest request); + + Project unbindRepository(Long workspaceId, Long projectId); + + void updateRepositorySettings(Long workspaceId, Long projectId, UpdateRepositorySettingsRequest request) + throws GeneralSecurityException; + + Project changeVcsConnection(Long workspaceId, Long projectId, ChangeVcsConnectionRequest request); + + // ==================== AI Connection ==================== + + boolean bindAiConnection(Long workspaceId, Long projectId, BindAiConnectionRequest request); + + // ==================== Branch Management ==================== + + List getProjectBranches(Long workspaceId, String namespace); + + Project setDefaultBranch(Long workspaceId, String namespace, Long branchId); + + Project setDefaultBranchByName(Long workspaceId, String namespace, String branchName); + + // ==================== Configuration ==================== + + BranchAnalysisConfig getBranchAnalysisConfig(Project project); + + Project updateBranchAnalysisConfig(Long workspaceId, Long projectId, + List prTargetBranches, List branchPushPatterns); + + Project updateRagConfig(Long workspaceId, Long projectId, boolean enabled, String branch, + List excludePatterns, Boolean multiBranchEnabled, Integer branchRetentionDays); + + Project updateRagConfig(Long workspaceId, Long projectId, boolean enabled, String branch, + List excludePatterns); + + Project updateAnalysisSettings(Long workspaceId, Long projectId, Boolean prAnalysisEnabled, + Boolean branchAnalysisEnabled, InstallationMethod installationMethod, Integer maxAnalysisTokenLimit); + + Project updateProjectQualityGate(Long workspaceId, Long projectId, Long qualityGateId); + + CommentCommandsConfig getCommentCommandsConfig(Project project); + + Project updateCommentCommandsConfig(Long workspaceId, Long projectId, UpdateCommentCommandsConfigRequest request); + + // ==================== Webhooks ==================== + + WebhookSetupResult setupWebhooks(Long workspaceId, Long projectId); + + WebhookInfo getWebhookInfo(Long workspaceId, Long projectId); + + // ==================== DTOs ==================== + + record WebhookSetupResult( + boolean success, + String webhookId, + String webhookUrl, + String message) { + } + + record WebhookInfo( + boolean webhooksConfigured, + String webhookId, + String webhookUrl, + EVcsProvider provider) { + } +} diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/ProjectService.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/ProjectService.java index 07aeb2be..aa2498cc 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/ProjectService.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/project/service/ProjectService.java @@ -58,9 +58,9 @@ import org.slf4j.LoggerFactory; @Service -public class ProjectService { +public class ProjectService implements IProjectService { private static final Logger log = LoggerFactory.getLogger(ProjectService.class); - + private final ProjectRepository projectRepository; private final VcsConnectionRepository vcsConnectionRepository; private final TokenEncryptionService tokenEncryptionService; @@ -103,8 +103,7 @@ public ProjectService( JobLogRepository jobLogRepository, PrSummarizeCacheRepository prSummarizeCacheRepository, VcsClientProvider vcsClientProvider, - QualityGateRepository qualityGateRepository - ) { + QualityGateRepository qualityGateRepository) { this.projectRepository = projectRepository; this.vcsConnectionRepository = vcsConnectionRepository; this.tokenEncryptionService = tokenEncryptionService; @@ -128,21 +127,21 @@ public ProjectService( @Transactional(readOnly = true) public List listWorkspaceProjects(Long workspaceId) { - // Use the method that fetches default branch eagerly to include stats in project list + // Use the method that fetches default branch eagerly to include stats in + // project list return projectRepository.findByWorkspaceIdWithDefaultBranch(workspaceId); } @Transactional(readOnly = true) public org.springframework.data.domain.Page listWorkspaceProjectsPaginated( - Long workspaceId, - String search, - int page, + Long workspaceId, + String search, + int page, int size) { org.springframework.data.domain.Pageable pageable = org.springframework.data.domain.PageRequest.of( - page, - size, - org.springframework.data.domain.Sort.by(org.springframework.data.domain.Sort.Direction.DESC, "id") - ); + page, + size, + org.springframework.data.domain.Sort.by(org.springframework.data.domain.Sort.Direction.DESC, "id")); return projectRepository.findByWorkspaceIdWithSearchPaginated(workspaceId, search, pageable); } @@ -159,7 +158,9 @@ public Project createProject(Long workspaceId, CreateProjectRequest request) thr // ensure namespace uniqueness per workspace projectRepository.findByWorkspaceIdAndNamespace(workspaceId, request.getNamespace()) - .ifPresent(p -> { throw new InvalidProjectRequestException("Project namespace already exists in workspace"); }); + .ifPresent(p -> { + throw new InvalidProjectRequestException("Project namespace already exists in workspace"); + }); Project newProject = new Project(); newProject.setWorkspace(ws); @@ -179,10 +180,12 @@ public Project createProject(Long workspaceId, CreateProjectRequest request) thr newProject.setConfiguration(config); if (request.hasVcsConnection()) { - VcsConnection vcsConnection = vcsConnectionRepository.findByWorkspace_IdAndId(workspaceId, request.getVcsConnectionId()) + VcsConnection vcsConnection = vcsConnectionRepository + .findByWorkspace_IdAndId(workspaceId, request.getVcsConnectionId()) .orElseThrow(() -> new NoSuchElementException("VCS connection not found!")); - // Use VcsRepoBinding (provider-agnostic) instead of legacy ProjectVcsConnectionBinding + // Use VcsRepoBinding (provider-agnostic) instead of legacy + // ProjectVcsConnectionBinding VcsRepoBinding vcsRepoBinding = new VcsRepoBinding(); vcsRepoBinding.setProject(newProject); vcsRepoBinding.setWorkspace(ws); @@ -193,7 +196,8 @@ public Project createProject(Long workspaceId, CreateProjectRequest request) thr if (request.getRepositoryUUID() != null) { vcsRepoBinding.setExternalRepoId(request.getRepositoryUUID().toString()); } - // For Bitbucket, extract workspace from config; for other providers, use a default + // For Bitbucket, extract workspace from config; for other providers, use a + // default String externalNamespace = null; if (vcsConnection.getConfiguration() instanceof BitbucketCloudConfig bbConfig) { externalNamespace = bbConfig.workspaceId(); @@ -206,7 +210,8 @@ public Project createProject(Long workspaceId, CreateProjectRequest request) thr } if (request.getAiConnectionId() != null) { - AIConnection aiConnection = aiConnectionRepository.findByWorkspace_IdAndId(workspaceId, request.getAiConnectionId()) + AIConnection aiConnection = aiConnectionRepository + .findByWorkspace_IdAndId(workspaceId, request.getAiConnectionId()) .orElseThrow(() -> new NoSuchElementException("AI connection not found!")); ProjectAiConnectionBinding aiBinding = new ProjectAiConnectionBinding(); @@ -222,7 +227,8 @@ public Project createProject(Long workspaceId, CreateProjectRequest request) thr String plainToken = Base64.getUrlEncoder().withoutPadding().encodeToString(random); String encrypted = tokenEncryptionService.encrypt(plainToken); newProject.setAuthToken(encrypted); - // Note: plainToken is not returned in API responses; store encrypted token only. + // Note: plainToken is not returned in API responses; store encrypted token + // only. } catch (GeneralSecurityException e) { throw new SecurityException("Failed to generate project auth token"); } @@ -345,7 +351,7 @@ public Project bindRepository(Long workspaceId, Long projectId, BindRepositoryRe VcsConnection conn = vcsConnectionRepository.findByWorkspace_IdAndId(workspaceId, request.getConnectionId()) .orElseThrow(() -> new NoSuchElementException("Connection not found")); } - //TODO: bind implementation + // TODO: bind implementation return projectRepository.save(p); } @@ -353,14 +359,14 @@ public Project bindRepository(Long workspaceId, Long projectId, BindRepositoryRe public Project unbindRepository(Long workspaceId, Long projectId) { Project p = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); - //TODO: unbind implementation + // TODO: unbind implementation // Clear settings placeholder if used in future return projectRepository.save(p); } - @Transactional - public void updateRepositorySettings(Long workspaceId, Long projectId, UpdateRepositorySettingsRequest request) throws GeneralSecurityException { + public void updateRepositorySettings(Long workspaceId, Long projectId, UpdateRepositorySettingsRequest request) + throws GeneralSecurityException { Project p = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); @@ -368,25 +374,42 @@ public void updateRepositorySettings(Long workspaceId, Long projectId, UpdateRep String encrypted = tokenEncryptionService.encrypt(request.getToken()); p.setAuthToken(encrypted); } - //TODO: service implementation + // TODO: service implementation // apiBaseUrl and webhookSecret ignored for now (no fields yet) projectRepository.save(p); } @Transactional - public boolean bindAiConnection(Long workspaceId, Long projectId, BindAiConnectionRequest request) throws SecurityException { - Project project = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) + public boolean bindAiConnection(Long workspaceId, Long projectId, BindAiConnectionRequest request) + throws SecurityException { + // Use findByIdWithConnections to eagerly fetch aiBinding for proper orphan + // removal + Project project = projectRepository.findByIdWithConnections(projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); + // Verify workspace ownership + if (!project.getWorkspace().getId().equals(workspaceId)) { + throw new NoSuchElementException("Project not found in workspace"); + } + if (request.getAiConnectionId() != null) { Long aiConnectionId = request.getAiConnectionId(); AIConnection aiConnection = aiConnectionRepository.findByWorkspace_IdAndId(workspaceId, aiConnectionId) .orElseThrow(() -> new NoSuchElementException("Ai connection not found")); - ProjectAiConnectionBinding aiConnectionBinding = new ProjectAiConnectionBinding(); - aiConnectionBinding.setProject(project); - aiConnectionBinding.setAiConnection(aiConnection); - project.setAiConnectionBinding(aiConnectionBinding); + // Check if there's an existing binding that needs to be updated + ProjectAiConnectionBinding existingBinding = project.getAiBinding(); + if (existingBinding != null) { + // Update existing binding instead of creating new one + existingBinding.setAiConnection(aiConnection); + } else { + // Create new binding + ProjectAiConnectionBinding aiConnectionBinding = new ProjectAiConnectionBinding(); + aiConnectionBinding.setProject(project); + aiConnectionBinding.setAiConnection(aiConnection); + project.setAiConnectionBinding(aiConnectionBinding); + } + projectRepository.save(project); return true; } @@ -446,7 +469,8 @@ public Project setDefaultBranchByName(Long workspaceId, String namespace, String Project project = getProjectByWorkspaceAndNamespace(workspaceId, namespace); Branch branch = branchRepository.findByProjectIdAndBranchName(project.getId(), branchName) - .orElseThrow(() -> new NoSuchElementException("Branch '" + branchName + "' not found for this project")); + .orElseThrow( + () -> new NoSuchElementException("Branch '" + branchName + "' not found for this project")); project.setDefaultBranch(branch); return projectRepository.save(project); @@ -467,10 +491,13 @@ public BranchAnalysisConfig getBranchAnalysisConfig(Project project) { /** * Update the branch analysis configuration for a project. * Main branch is always ensured to be in the patterns. - * @param workspaceId the workspace ID - * @param projectId the project ID - * @param prTargetBranches patterns for PR target branches (e.g., ["main", "develop", "release/*"]) - * @param branchPushPatterns patterns for branch push analysis (e.g., ["main", "develop"]) + * + * @param workspaceId the workspace ID + * @param projectId the project ID + * @param prTargetBranches patterns for PR target branches (e.g., ["main", + * "develop", "release/*"]) + * @param branchPushPatterns patterns for branch push analysis (e.g., ["main", + * "develop"]) * @return the updated project */ @Transactional @@ -478,8 +505,7 @@ public Project updateBranchAnalysisConfig( Long workspaceId, Long projectId, List prTargetBranches, - List branchPushPatterns - ) { + List branchPushPatterns) { Project project = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); @@ -490,8 +516,7 @@ public Project updateBranchAnalysisConfig( BranchAnalysisConfig branchConfig = new BranchAnalysisConfig( prTargetBranches, - branchPushPatterns - ); + branchPushPatterns); currentConfig.setBranchAnalysis(branchConfig); // Ensure main branch is always included in patterns @@ -502,12 +527,14 @@ public Project updateBranchAnalysisConfig( /** * Update the RAG configuration for a project. - * @param workspaceId the workspace ID - * @param projectId the project ID - * @param enabled whether RAG indexing is enabled - * @param branch the branch to index (null uses defaultBranch or 'main') - * @param excludePatterns patterns to exclude from indexing - * @param multiBranchEnabled whether multi-branch indexing is enabled + * + * @param workspaceId the workspace ID + * @param projectId the project ID + * @param enabled whether RAG indexing is enabled + * @param branch the branch to index (null uses defaultBranch or + * 'main') + * @param excludePatterns patterns to exclude from indexing + * @param multiBranchEnabled whether multi-branch indexing is enabled * @param branchRetentionDays how long to keep branch index metadata * @return the updated project */ @@ -519,8 +546,7 @@ public Project updateRagConfig( String branch, java.util.List excludePatterns, Boolean multiBranchEnabled, - Integer branchRetentionDays - ) { + Integer branchRetentionDays) { Project project = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); @@ -540,7 +566,7 @@ public Project updateRagConfig( prAnalysisEnabled, branchAnalysisEnabled, installationMethod, commentCommands)); return projectRepository.save(project); } - + /** * Simplified RAG config update (backward compatible). */ @@ -550,8 +576,7 @@ public Project updateRagConfig( Long projectId, boolean enabled, String branch, - java.util.List excludePatterns - ) { + java.util.List excludePatterns) { return updateRagConfig(workspaceId, projectId, enabled, branch, excludePatterns, null, null); } @@ -561,8 +586,8 @@ public Project updateAnalysisSettings( Long projectId, Boolean prAnalysisEnabled, Boolean branchAnalysisEnabled, - InstallationMethod installationMethod - ) { + InstallationMethod installationMethod, + Integer maxAnalysisTokenLimit) { Project project = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); @@ -572,28 +597,32 @@ public Project updateAnalysisSettings( var branchAnalysis = currentConfig != null ? currentConfig.branchAnalysis() : null; var ragConfig = currentConfig != null ? currentConfig.ragConfig() : null; var commentCommands = currentConfig != null ? currentConfig.commentCommands() : null; + int currentMaxTokenLimit = currentConfig != null ? currentConfig.maxAnalysisTokenLimit() + : ProjectConfig.DEFAULT_MAX_ANALYSIS_TOKEN_LIMIT; - Boolean newPrAnalysis = prAnalysisEnabled != null ? prAnalysisEnabled : - (currentConfig != null ? currentConfig.prAnalysisEnabled() : true); - Boolean newBranchAnalysis = branchAnalysisEnabled != null ? branchAnalysisEnabled : - (currentConfig != null ? currentConfig.branchAnalysisEnabled() : true); - var newInstallationMethod = installationMethod != null ? installationMethod : - (currentConfig != null ? currentConfig.installationMethod() : null); + Boolean newPrAnalysis = prAnalysisEnabled != null ? prAnalysisEnabled + : (currentConfig != null ? currentConfig.prAnalysisEnabled() : true); + Boolean newBranchAnalysis = branchAnalysisEnabled != null ? branchAnalysisEnabled + : (currentConfig != null ? currentConfig.branchAnalysisEnabled() : true); + var newInstallationMethod = installationMethod != null ? installationMethod + : (currentConfig != null ? currentConfig.installationMethod() : null); + int newMaxTokenLimit = maxAnalysisTokenLimit != null ? maxAnalysisTokenLimit : currentMaxTokenLimit; // Update both the direct column and the JSON config - //TODO: remove duplication + // TODO: remove duplication project.setPrAnalysisEnabled(newPrAnalysis != null ? newPrAnalysis : true); project.setBranchAnalysisEnabled(newBranchAnalysis != null ? newBranchAnalysis : true); project.setConfiguration(new ProjectConfig(useLocalMcp, mainBranch, branchAnalysis, ragConfig, - newPrAnalysis, newBranchAnalysis, newInstallationMethod, commentCommands)); + newPrAnalysis, newBranchAnalysis, newInstallationMethod, commentCommands, newMaxTokenLimit)); return projectRepository.save(project); } /** * Update the quality gate for a project. - * @param workspaceId the workspace ID - * @param projectId the project ID + * + * @param workspaceId the workspace ID + * @param projectId the project ID * @param qualityGateId the quality gate ID (null to remove) * @return the updated project */ @@ -601,7 +630,7 @@ public Project updateAnalysisSettings( public Project updateProjectQualityGate(Long workspaceId, Long projectId, Long qualityGateId) { Project project = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); - + if (qualityGateId != null) { QualityGate qualityGate = qualityGateRepository.findByIdAndWorkspaceId(qualityGateId, workspaceId) .orElseThrow(() -> new NoSuchElementException("Quality gate not found")); @@ -609,13 +638,14 @@ public Project updateProjectQualityGate(Long workspaceId, Long projectId, Long q } else { project.setQualityGate(null); } - + return projectRepository.save(project); } /** * Get the comment commands configuration for a project. - * Returns a CommentCommandsConfig record (never null, returns default disabled config if not configured). + * Returns a CommentCommandsConfig record (never null, returns default disabled + * config if not configured). */ @Transactional(readOnly = true) public CommentCommandsConfig getCommentCommandsConfig(Project project) { @@ -627,17 +657,17 @@ public CommentCommandsConfig getCommentCommandsConfig(Project project) { /** * Update the comment commands configuration for a project. + * * @param workspaceId the workspace ID - * @param projectId the project ID - * @param request the update request + * @param projectId the project ID + * @param request the update request * @return the updated project */ @Transactional public Project updateCommentCommandsConfig( Long workspaceId, Long projectId, - org.rostilos.codecrow.webserver.project.dto.request.UpdateCommentCommandsConfigRequest request - ) { + org.rostilos.codecrow.webserver.project.dto.request.UpdateCommentCommandsConfigRequest request) { Project project = projectRepository.findByWorkspaceIdAndId(workspaceId, projectId) .orElseThrow(() -> new NoSuchElementException("Project not found")); @@ -653,25 +683,27 @@ public Project updateCommentCommandsConfig( // Build new comment commands config var existingCommentConfig = currentConfig != null ? currentConfig.commentCommands() : null; - boolean enabled = request.enabled() != null ? request.enabled() : - (existingCommentConfig != null ? existingCommentConfig.enabled() : false); - Integer rateLimit = request.rateLimit() != null ? request.rateLimit() : - (existingCommentConfig != null ? existingCommentConfig.rateLimit() : CommentCommandsConfig.DEFAULT_RATE_LIMIT); - Integer rateLimitWindow = request.rateLimitWindowMinutes() != null ? request.rateLimitWindowMinutes() : - (existingCommentConfig != null ? existingCommentConfig.rateLimitWindowMinutes() : CommentCommandsConfig.DEFAULT_RATE_LIMIT_WINDOW_MINUTES); - Boolean allowPublicRepoCommands = request.allowPublicRepoCommands() != null ? request.allowPublicRepoCommands() : - (existingCommentConfig != null ? existingCommentConfig.allowPublicRepoCommands() : false); - List allowedCommands = request.allowedCommands() != null ? request.validatedAllowedCommands() : - (existingCommentConfig != null ? existingCommentConfig.allowedCommands() : null); - CommandAuthorizationMode authorizationMode = request.authorizationMode() != null ? request.authorizationMode() : - (existingCommentConfig != null ? existingCommentConfig.authorizationMode() : CommentCommandsConfig.DEFAULT_AUTHORIZATION_MODE); - Boolean allowPrAuthor = request.allowPrAuthor() != null ? request.allowPrAuthor() : - (existingCommentConfig != null ? existingCommentConfig.allowPrAuthor() : true); + boolean enabled = request.enabled() != null ? request.enabled() + : (existingCommentConfig != null ? existingCommentConfig.enabled() : false); + Integer rateLimit = request.rateLimit() != null ? request.rateLimit() + : (existingCommentConfig != null ? existingCommentConfig.rateLimit() + : CommentCommandsConfig.DEFAULT_RATE_LIMIT); + Integer rateLimitWindow = request.rateLimitWindowMinutes() != null ? request.rateLimitWindowMinutes() + : (existingCommentConfig != null ? existingCommentConfig.rateLimitWindowMinutes() + : CommentCommandsConfig.DEFAULT_RATE_LIMIT_WINDOW_MINUTES); + Boolean allowPublicRepoCommands = request.allowPublicRepoCommands() != null ? request.allowPublicRepoCommands() + : (existingCommentConfig != null ? existingCommentConfig.allowPublicRepoCommands() : false); + List allowedCommands = request.allowedCommands() != null ? request.validatedAllowedCommands() + : (existingCommentConfig != null ? existingCommentConfig.allowedCommands() : null); + CommandAuthorizationMode authorizationMode = request.authorizationMode() != null ? request.authorizationMode() + : (existingCommentConfig != null ? existingCommentConfig.authorizationMode() + : CommentCommandsConfig.DEFAULT_AUTHORIZATION_MODE); + Boolean allowPrAuthor = request.allowPrAuthor() != null ? request.allowPrAuthor() + : (existingCommentConfig != null ? existingCommentConfig.allowPrAuthor() : true); var commentCommands = new CommentCommandsConfig( enabled, rateLimit, rateLimitWindow, allowPublicRepoCommands, allowedCommands, - authorizationMode, allowPrAuthor - ); + authorizationMode, allowPrAuthor); project.setConfiguration(new ProjectConfig(useLocalMcp, mainBranch, branchAnalysis, ragConfig, prAnalysisEnabled, branchAnalysisEnabled, installationMethod, commentCommands)); @@ -720,14 +752,14 @@ public WebhookSetupResult setupWebhooks(Long workspaceId, Long projectId) { // Get VCS client and setup webhook org.rostilos.codecrow.vcsclient.VcsClient client = vcsClientProvider.getClient(connection); List events = getWebhookEvents(binding.getProvider()); - + String workspaceIdOrNamespace; String repoSlug; - + // For REPOSITORY_TOKEN connections, use the repositoryPath from the connection // because the token is scoped to that specific repository - if (connection.getConnectionType() == EVcsConnectionType.REPOSITORY_TOKEN - && connection.getRepositoryPath() != null + if (connection.getConnectionType() == EVcsConnectionType.REPOSITORY_TOKEN + && connection.getRepositoryPath() != null && !connection.getRepositoryPath().isBlank()) { String repositoryPath = connection.getRepositoryPath(); int lastSlash = repositoryPath.lastIndexOf('/'); @@ -738,16 +770,16 @@ public WebhookSetupResult setupWebhooks(Long workspaceId, Long projectId) { workspaceIdOrNamespace = binding.getExternalNamespace(); repoSlug = repositoryPath; } - log.info("REPOSITORY_TOKEN webhook setup - using repositoryPath: {}, namespace: {}, slug: {}", + log.info("REPOSITORY_TOKEN webhook setup - using repositoryPath: {}, namespace: {}, slug: {}", repositoryPath, workspaceIdOrNamespace, repoSlug); } else { workspaceIdOrNamespace = binding.getExternalNamespace(); repoSlug = binding.getExternalRepoSlug(); log.info("Standard webhook setup - namespace: {}, slug: {}", workspaceIdOrNamespace, repoSlug); } - + String webhookId = client.ensureWebhook(workspaceIdOrNamespace, repoSlug, webhookUrl, events); - + if (webhookId != null) { binding.setWebhookId(webhookId); binding.setWebhooksConfigured(true); @@ -781,8 +813,7 @@ public WebhookInfo getWebhookInfo(Long workspaceId, Long projectId) { binding.isWebhooksConfigured(), binding.getWebhookId(), webhookUrl, - binding.getProvider() - ); + binding.getProvider()); } private String generateWebhookUrl(EVcsProvider provider, Project project) { @@ -792,7 +823,8 @@ private String generateWebhookUrl(EVcsProvider provider, Project project) { private List getWebhookEvents(EVcsProvider provider) { return switch (provider) { - case BITBUCKET_CLOUD -> List.of("pullrequest:created", "pullrequest:updated", "pullrequest:fulfilled", "pullrequest:comment_created", "repo:push"); + case BITBUCKET_CLOUD -> List.of("pullrequest:created", "pullrequest:updated", "pullrequest:fulfilled", + "pullrequest:comment_created", "repo:push"); case GITHUB -> List.of("pull_request", "pull_request_review_comment", "issue_comment", "push"); case GITLAB -> List.of("merge_requests_events", "note_events", "push_events"); default -> List.of(); @@ -816,7 +848,8 @@ public Project changeVcsConnection(Long workspaceId, Long projectId, ChangeVcsCo Workspace workspace = workspaceRepository.findById(workspaceId) .orElseThrow(() -> new NoSuchElementException("Workspace not found")); - VcsConnection newConnection = vcsConnectionRepository.findByWorkspace_IdAndId(workspaceId, request.getConnectionId()) + VcsConnection newConnection = vcsConnectionRepository + .findByWorkspace_IdAndId(workspaceId, request.getConnectionId()) .orElseThrow(() -> new NoSuchElementException("VCS Connection not found")); // Get or create VcsRepoBinding @@ -839,9 +872,11 @@ public Project changeVcsConnection(Long workspaceId, Long projectId, ChangeVcsCo binding.setVcsConnection(newConnection); binding.setProvider(newConnection.getProviderType()); binding.setExternalRepoSlug(request.getRepositorySlug()); - binding.setExternalNamespace(request.getWorkspaceId() != null ? request.getWorkspaceId() : newConnection.getExternalWorkspaceSlug()); - binding.setExternalRepoId(request.getRepositoryId() != null ? request.getRepositoryId() : request.getRepositorySlug()); - + binding.setExternalNamespace( + request.getWorkspaceId() != null ? request.getWorkspaceId() : newConnection.getExternalWorkspaceSlug()); + binding.setExternalRepoId( + request.getRepositoryId() != null ? request.getRepositoryId() : request.getRepositorySlug()); + if (request.getDefaultBranch() != null && !request.getDefaultBranch().isBlank()) { binding.setDefaultBranch(request.getDefaultBranch()); } @@ -882,20 +917,4 @@ private void clearProjectAnalysisData(Long projectId) { analysisLockRepository.deleteByProjectId(projectId); ragIndexStatusRepository.deleteByProjectId(projectId); } - - // ==================== DTOs for Webhook Operations ==================== - - public record WebhookSetupResult( - boolean success, - String webhookId, - String webhookUrl, - String message - ) {} - - public record WebhookInfo( - boolean webhooksConfigured, - String webhookId, - String webhookUrl, - EVcsProvider provider - ) {} } diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/qualitygate/controller/QualityGateController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/qualitygate/controller/QualityGateController.java index 7e47cef1..b8c7d3de 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/qualitygate/controller/QualityGateController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/qualitygate/controller/QualityGateController.java @@ -6,13 +6,14 @@ import org.rostilos.codecrow.core.dto.qualitygate.QualityGateDTO; import org.rostilos.codecrow.core.model.qualitygate.QualityGate; import org.rostilos.codecrow.core.model.workspace.Workspace; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; import org.rostilos.codecrow.webserver.qualitygate.dto.request.CreateQualityGateRequest; import org.rostilos.codecrow.webserver.qualitygate.dto.request.UpdateQualityGateRequest; import org.rostilos.codecrow.webserver.qualitygate.service.QualityGateService; import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.*; import java.util.List; @@ -34,7 +35,7 @@ public QualityGateController( } @GetMapping - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity> listQualityGates(@PathVariable String workspaceSlug) { Workspace workspace = workspaceService.getWorkspaceBySlug(workspaceSlug); // Ensure a default quality gate exists, creating one if needed @@ -53,7 +54,7 @@ public ResponseEntity> listQualityGates(@PathVariable Strin } @GetMapping("/{qualityGateId}") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity getQualityGate( @PathVariable String workspaceSlug, @PathVariable Long qualityGateId @@ -64,7 +65,7 @@ public ResponseEntity getQualityGate( } @GetMapping("/default") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity getDefaultQualityGate(@PathVariable String workspaceSlug) { Workspace workspace = workspaceService.getWorkspaceBySlug(workspaceSlug); // Ensure a default quality gate exists, creating one if needed @@ -82,7 +83,7 @@ public ResponseEntity getDefaultQualityGate(@PathVariable String } @PostMapping - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity createQualityGate( @PathVariable String workspaceSlug, @Valid @RequestBody CreateQualityGateRequest request @@ -93,7 +94,7 @@ public ResponseEntity createQualityGate( } @PutMapping("/{qualityGateId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity updateQualityGate( @PathVariable String workspaceSlug, @PathVariable Long qualityGateId, @@ -105,7 +106,7 @@ public ResponseEntity updateQualityGate( } @PostMapping("/{qualityGateId}/set-default") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity setDefault( @PathVariable String workspaceSlug, @PathVariable Long qualityGateId @@ -116,7 +117,7 @@ public ResponseEntity setDefault( } @DeleteMapping("/{qualityGateId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity deleteQualityGate( @PathVariable String workspaceSlug, @PathVariable Long qualityGateId diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/cloud/BitbucketCloudController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/cloud/BitbucketCloudController.java index 5c0d4027..fdfd55f0 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/cloud/BitbucketCloudController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/cloud/BitbucketCloudController.java @@ -7,6 +7,7 @@ import org.rostilos.codecrow.core.model.vcs.EVcsProvider; import org.rostilos.codecrow.core.model.vcs.VcsConnection; import org.rostilos.codecrow.core.persistence.repository.vcs.VcsConnectionRepository; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; import org.rostilos.codecrow.vcsclient.bitbucket.cloud.dto.response.RepositorySearchResult; import org.rostilos.codecrow.webserver.vcs.dto.request.cloud.BitbucketCloudCreateRequest; import org.rostilos.codecrow.core.dto.bitbucket.BitbucketCloudDTO; @@ -15,7 +16,6 @@ import org.rostilos.codecrow.webserver.vcs.utils.BitbucketCloudConfigHandler; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.CrossOrigin; import org.springframework.web.bind.annotation.DeleteMapping; import org.springframework.web.bind.annotation.GetMapping; @@ -29,6 +29,7 @@ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@HasOwnerOrAdminRights @RequestMapping("/api/{workspaceSlug}/vcs/bitbucket_cloud") public class BitbucketCloudController { private final VcsConnectionRepository vcsConnectionRepository; @@ -49,7 +50,6 @@ public BitbucketCloudController( } @GetMapping("/list") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity> getUserBitbucketCloudConnections(@PathVariable String workspaceSlug) { Long workspaceId = workspaceService.getWorkspaceBySlug(workspaceSlug).getId(); List userBitbucketConnections = vcsConnectionService.getWorkspaceBitbucketCloudConnections(workspaceId); @@ -61,7 +61,6 @@ public ResponseEntity> getUserBitbucketCloudConnections( } @PostMapping("/create") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity createUserBitbucketCloudConnection( @PathVariable String workspaceSlug, @RequestBody BitbucketCloudCreateRequest request @@ -78,7 +77,6 @@ public ResponseEntity createUserBitbucketCloudConnection( } @PatchMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity updateUserBitbucketCloudConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId, @@ -94,7 +92,6 @@ public ResponseEntity updateUserBitbucketCloudConnection( } @DeleteMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity deleteUserBitbucketCloudConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId @@ -105,7 +102,6 @@ public ResponseEntity deleteUserBitbucketCloudConnection( } @GetMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getConnectionInfo( @PathVariable String workspaceSlug, @PathVariable Long connectionId @@ -121,7 +117,6 @@ public ResponseEntity getConnectionInfo( @GetMapping("/{connectionId}/repositories") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getBitbucketConnectionRepositories( @PathVariable String workspaceSlug, @PathVariable Long connectionId, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/github/GitHubController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/github/GitHubController.java index b6060b6e..fa21bc03 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/github/GitHubController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/github/GitHubController.java @@ -8,13 +8,13 @@ import org.rostilos.codecrow.core.model.vcs.VcsConnection; import org.rostilos.codecrow.core.model.vcs.config.github.GitHubConfig; import org.rostilos.codecrow.core.persistence.repository.vcs.VcsConnectionRepository; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; import org.rostilos.codecrow.vcsclient.github.dto.response.RepositorySearchResult; import org.rostilos.codecrow.webserver.vcs.dto.request.github.GitHubCreateRequest; import org.rostilos.codecrow.webserver.vcs.service.VcsConnectionWebService; import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.CrossOrigin; import org.springframework.web.bind.annotation.DeleteMapping; import org.springframework.web.bind.annotation.GetMapping; @@ -28,6 +28,7 @@ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@HasOwnerOrAdminRights @RequestMapping("/api/{workspaceSlug}/vcs/github") public class GitHubController { @@ -46,7 +47,6 @@ public GitHubController( } @GetMapping("/list") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity> getGitHubConnections(@PathVariable String workspaceSlug) { Long workspaceId = workspaceService.getWorkspaceBySlug(workspaceSlug).getId(); List connections = vcsConnectionService.getWorkspaceGitHubConnections(workspaceId); @@ -58,7 +58,6 @@ public ResponseEntity> getGitHubConnections(@PathVariable String } @PostMapping("/create") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity createGitHubConnection( @PathVariable String workspaceSlug, @RequestBody GitHubCreateRequest request @@ -81,7 +80,6 @@ public ResponseEntity createGitHubConnection( } @PatchMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity updateGitHubConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId, @@ -97,7 +95,6 @@ public ResponseEntity updateGitHubConnection( } @DeleteMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity deleteGitHubConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId @@ -108,7 +105,6 @@ public ResponseEntity deleteGitHubConnection( } @GetMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getConnectionInfo( @PathVariable String workspaceSlug, @PathVariable Long connectionId @@ -125,7 +121,6 @@ public ResponseEntity getConnectionInfo( } @GetMapping("/{connectionId}/repositories") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getGitHubConnectionRepositories( @PathVariable String workspaceSlug, @PathVariable Long connectionId, diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/gitlab/GitLabController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/gitlab/GitLabController.java index e6d469a3..e6f3430c 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/gitlab/GitLabController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/vcs/controller/gitlab/GitLabController.java @@ -8,6 +8,7 @@ import org.rostilos.codecrow.core.model.vcs.VcsConnection; import org.rostilos.codecrow.core.model.vcs.config.gitlab.GitLabConfig; import org.rostilos.codecrow.core.persistence.repository.vcs.VcsConnectionRepository; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; import org.rostilos.codecrow.vcsclient.gitlab.dto.response.RepositorySearchResult; import org.rostilos.codecrow.webserver.vcs.dto.request.gitlab.GitLabCreateRequest; import org.rostilos.codecrow.webserver.vcs.dto.request.gitlab.GitLabRepositoryTokenRequest; @@ -15,7 +16,6 @@ import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.web.bind.annotation.CrossOrigin; import org.springframework.web.bind.annotation.DeleteMapping; import org.springframework.web.bind.annotation.GetMapping; @@ -29,6 +29,7 @@ @CrossOrigin(origins = "*", maxAge = 3600) @RestController +@HasOwnerOrAdminRights @RequestMapping("/api/{workspaceSlug}/vcs/gitlab") public class GitLabController { @@ -47,7 +48,6 @@ public GitLabController( } @GetMapping("/list") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity> getGitLabConnections(@PathVariable String workspaceSlug) { Long workspaceId = workspaceService.getWorkspaceBySlug(workspaceSlug).getId(); List connections = vcsConnectionService.getWorkspaceGitLabConnections(workspaceId); @@ -59,7 +59,6 @@ public ResponseEntity> getGitLabConnections(@PathVariable String } @PostMapping("/create") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity createGitLabConnection( @PathVariable String workspaceSlug, @RequestBody GitLabCreateRequest request @@ -82,7 +81,6 @@ public ResponseEntity createGitLabConnection( } @PatchMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity updateGitLabConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId, @@ -98,7 +96,6 @@ public ResponseEntity updateGitLabConnection( } @DeleteMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity deleteGitLabConnection( @PathVariable String workspaceSlug, @PathVariable Long connectionId @@ -109,7 +106,6 @@ public ResponseEntity deleteGitLabConnection( } @GetMapping("/connections/{connectionId}") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getConnectionInfo( @PathVariable String workspaceSlug, @PathVariable Long connectionId @@ -126,7 +122,6 @@ public ResponseEntity getConnectionInfo( } @GetMapping("/{connectionId}/repositories") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity getGitLabConnectionRepositories( @PathVariable String workspaceSlug, @PathVariable Long connectionId, @@ -145,7 +140,6 @@ public ResponseEntity getGitLabConnectionRepositories( * Project Access Tokens are limited to a single project/repository. */ @PostMapping("/create-repository-token") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") public ResponseEntity createGitLabRepositoryTokenConnection( @PathVariable String workspaceSlug, @RequestBody GitLabRepositoryTokenRequest request diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceController.java index 97fdd275..1a537a57 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceController.java @@ -8,6 +8,9 @@ import org.rostilos.codecrow.core.model.workspace.EWorkspaceRole; import org.rostilos.codecrow.core.model.workspace.Workspace; import org.rostilos.codecrow.core.model.workspace.WorkspaceMember; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; +import org.rostilos.codecrow.security.annotations.IsWorkspaceOwner; import org.rostilos.codecrow.security.service.UserDetailsImpl; import org.rostilos.codecrow.webserver.generic.dto.message.MessageResponse; import org.rostilos.codecrow.webserver.workspace.dto.request.ChangeRoleRequest; @@ -60,7 +63,7 @@ public ResponseEntity createWorkspace( } @PostMapping("/{workspaceSlug}/invite") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity inviteToWorkspace( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -73,7 +76,7 @@ public ResponseEntity inviteToWorkspace( } @DeleteMapping("/{workspaceSlug}/member/remove") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity removeMember( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -85,7 +88,7 @@ public ResponseEntity removeMember( } @PutMapping("/{workspaceSlug}/changeRole") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity changeWorkspaceRole( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -111,7 +114,7 @@ public ResponseEntity acceptInvite( } @GetMapping("/{workspaceSlug}/members") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> listMembers( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug @@ -123,7 +126,7 @@ public ResponseEntity> listMembers( } @GetMapping("/{workspaceSlug}/role") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity getUserRole( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug @@ -136,7 +139,7 @@ public ResponseEntity getUserRole( } @DeleteMapping("/{workspaceSlug}") - @PreAuthorize("@workspaceSecurity.isWorkspaceOwner(#workspaceSlug, authentication)") + @IsWorkspaceOwner public ResponseEntity deleteWorkspace( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -155,7 +158,7 @@ public ResponseEntity deleteWorkspace( } @PostMapping("/{workspaceSlug}/cancel-deletion") - @PreAuthorize("@workspaceSecurity.isWorkspaceOwner(#workspaceSlug, authentication)") + @IsWorkspaceOwner public ResponseEntity cancelScheduledDeletion( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug @@ -165,7 +168,7 @@ public ResponseEntity cancelScheduledDeletion( } @GetMapping("/{workspaceSlug}/deletion-status") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity getDeletionStatus( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceOwnershipTransferController.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceOwnershipTransferController.java index cabe834d..5e137639 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceOwnershipTransferController.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/controller/WorkspaceOwnershipTransferController.java @@ -2,6 +2,9 @@ import jakarta.validation.Valid; import org.rostilos.codecrow.core.model.workspace.WorkspaceOwnershipTransfer; +import org.rostilos.codecrow.security.annotations.HasOwnerOrAdminRights; +import org.rostilos.codecrow.security.annotations.IsWorkspaceMember; +import org.rostilos.codecrow.security.annotations.IsWorkspaceOwner; import org.rostilos.codecrow.security.service.UserDetailsImpl; import org.rostilos.codecrow.webserver.generic.dto.message.MessageResponse; import org.rostilos.codecrow.webserver.workspace.dto.request.CancelOwnershipTransferRequest; @@ -11,7 +14,6 @@ import org.rostilos.codecrow.webserver.workspace.service.WorkspaceService; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; -import org.springframework.security.access.prepost.PreAuthorize; import org.springframework.security.core.annotation.AuthenticationPrincipal; import org.springframework.web.bind.annotation.*; @@ -38,7 +40,7 @@ public WorkspaceOwnershipTransferController( * Initiate ownership transfer. Requires 2FA from the current owner. */ @PostMapping("/{workspaceSlug}/ownership/transfer") - @PreAuthorize("@workspaceSecurity.isWorkspaceOwner(#workspaceSlug, authentication)") + @IsWorkspaceOwner public ResponseEntity initiateTransfer( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -58,7 +60,7 @@ public ResponseEntity initiateTransfer( * Cancel a pending ownership transfer. */ @DeleteMapping("/{workspaceSlug}/ownership/transfer/{transferId}") - @PreAuthorize("@workspaceSecurity.isWorkspaceOwner(#workspaceSlug, authentication)") + @IsWorkspaceOwner public ResponseEntity cancelTransfer( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -74,7 +76,7 @@ public ResponseEntity cancelTransfer( * Complete an ownership transfer (by current owner or new owner). */ @PostMapping("/{workspaceSlug}/ownership/transfer/{transferId}/complete") - @PreAuthorize("@workspaceSecurity.isWorkspaceMember(#workspaceSlug, authentication)") + @IsWorkspaceMember public ResponseEntity completeTransfer( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug, @@ -88,7 +90,7 @@ public ResponseEntity completeTransfer( * Get pending transfer for a workspace, if any. */ @GetMapping("/{workspaceSlug}/ownership/transfer/pending") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity getPendingTransfer( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug @@ -103,7 +105,7 @@ public ResponseEntity getPendingTransfer( * Get transfer history for a workspace. */ @GetMapping("/{workspaceSlug}/ownership/transfer/history") - @PreAuthorize("@workspaceSecurity.hasOwnerOrAdminRights(#workspaceSlug, authentication)") + @HasOwnerOrAdminRights public ResponseEntity> getTransferHistory( @AuthenticationPrincipal UserDetailsImpl userDetails, @PathVariable String workspaceSlug diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/IWorkspaceService.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/IWorkspaceService.java new file mode 100644 index 00000000..b48331a9 --- /dev/null +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/IWorkspaceService.java @@ -0,0 +1,64 @@ +package org.rostilos.codecrow.webserver.workspace.service; + +import java.util.List; + +import org.rostilos.codecrow.core.model.workspace.EWorkspaceRole; +import org.rostilos.codecrow.core.model.workspace.Workspace; +import org.rostilos.codecrow.core.model.workspace.WorkspaceMember; + +/** + * Interface for workspace management operations. + *

+ * This interface enables cloud implementations to extend or decorate + * the base workspace service with additional capabilities like billing limits. + */ +public interface IWorkspaceService { + + // ==================== Core CRUD ==================== + + Workspace createWorkspace(Long actorUserId, String slug, String name, String description); + + Workspace getWorkspaceById(Long id); + + Workspace getWorkspaceBySlug(String slug); + + void deleteWorkspace(Long actorUserId, String workspaceSlug); + + // ==================== Membership ==================== + + WorkspaceMember inviteToWorkspace(Long actorUserId, Long workspaceId, String username, EWorkspaceRole role); + + WorkspaceMember inviteToWorkspace(Long actorUserId, String workspaceSlug, String username, EWorkspaceRole role); + + void removeMemberFromWorkspace(Long actorUserId, Long workspaceId, String username); + + void removeMemberFromWorkspace(Long actorUserId, String workspaceSlug, String username); + + void changeWorkspaceRole(Long actorUserId, Long workspaceId, String targetUsername, EWorkspaceRole newRole); + + void changeWorkspaceRole(Long actorUserId, String workspaceSlug, String targetUsername, EWorkspaceRole newRole); + + WorkspaceMember acceptInvite(Long userId, Long workspaceId); + + WorkspaceMember acceptInvite(Long userId, String workspaceSlug); + + // ==================== Queries ==================== + + List listMembers(Long workspaceId); + + List listMembers(String workspaceSlug); + + List listMembers(Long workspaceId, List excludeUsernames); + + List listUserWorkspaces(Long userId); + + EWorkspaceRole getUserRole(Long workspaceId, Long userId); + + EWorkspaceRole getUserRole(String workspaceSlug, Long userId); + + // ==================== Deletion Scheduling ==================== + + Workspace scheduleDeletion(Long actorUserId, String workspaceSlug); + + Workspace cancelScheduledDeletion(Long actorUserId, String workspaceSlug); +} diff --git a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/WorkspaceService.java b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/WorkspaceService.java index 0b991dab..65b134b3 100644 --- a/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/WorkspaceService.java +++ b/java-ecosystem/services/web-server/src/main/java/org/rostilos/codecrow/webserver/workspace/service/WorkspaceService.java @@ -19,7 +19,7 @@ import jakarta.persistence.PersistenceContext; @Service -public class WorkspaceService { +public class WorkspaceService implements IWorkspaceService { @PersistenceContext private EntityManager entityManager; @@ -29,8 +29,8 @@ public class WorkspaceService { private final UserRepository userRepository; public WorkspaceService(WorkspaceRepository workspaceRepository, - WorkspaceMemberRepository workspaceMemberRepository, - UserRepository userRepository) { + WorkspaceMemberRepository workspaceMemberRepository, + UserRepository userRepository) { this.workspaceRepository = workspaceRepository; this.workspaceMemberRepository = workspaceMemberRepository; this.userRepository = userRepository; @@ -105,7 +105,8 @@ public WorkspaceMember inviteToWorkspace(Long actorUserId, Long workspaceId, Str } @Transactional - public WorkspaceMember inviteToWorkspace(Long actorUserId, String workspaceSlug, String username, EWorkspaceRole role) { + public WorkspaceMember inviteToWorkspace(Long actorUserId, String workspaceSlug, String username, + EWorkspaceRole role) { Workspace workspace = getWorkspaceBySlug(workspaceSlug); return inviteToWorkspace(actorUserId, workspace.getId(), username, role); } @@ -115,16 +116,19 @@ public void changeWorkspaceRole(Long actorUserId, Long workspaceId, String targe User target = userRepository.findByUsername(targetUsername) .orElseThrow(() -> new IllegalArgumentException("User not found")); - Optional optionalWorkspaceMemberActor = workspaceMemberRepository.findByWorkspaceIdAndUserId(workspaceId, actorUserId); - if(optionalWorkspaceMemberActor.isEmpty()) { + Optional optionalWorkspaceMemberActor = workspaceMemberRepository + .findByWorkspaceIdAndUserId(workspaceId, actorUserId); + if (optionalWorkspaceMemberActor.isEmpty()) { throw new SecurityException("The current user does not have editing privileges."); } WorkspaceMember workspaceMemberActor = optionalWorkspaceMemberActor.get(); - if(workspaceMemberActor.getRole() == EWorkspaceRole.ADMIN && ( newRole == EWorkspaceRole.ADMIN || newRole == EWorkspaceRole.OWNER )) { + if (workspaceMemberActor.getRole() == EWorkspaceRole.ADMIN + && (newRole == EWorkspaceRole.ADMIN || newRole == EWorkspaceRole.OWNER)) { throw new SecurityException("The current user does not have editing privileges."); } - WorkspaceMember workspaceMember = workspaceMemberRepository.findByWorkspaceIdAndUserId(workspaceId, target.getId()) + WorkspaceMember workspaceMember = workspaceMemberRepository + .findByWorkspaceIdAndUserId(workspaceId, target.getId()) .orElseThrow(() -> new NoSuchElementException("The user is not present in the workspace.")); workspaceMember.setRole(newRole); @@ -132,7 +136,8 @@ public void changeWorkspaceRole(Long actorUserId, Long workspaceId, String targe } @Transactional - public void changeWorkspaceRole(Long actorUserId, String workspaceSlug, String targetUsername, EWorkspaceRole newRole) { + public void changeWorkspaceRole(Long actorUserId, String workspaceSlug, String targetUsername, + EWorkspaceRole newRole) { Workspace workspace = getWorkspaceBySlug(workspaceSlug); changeWorkspaceRole(actorUserId, workspace.getId(), targetUsername, newRole); } @@ -172,7 +177,8 @@ public List listMembers(String workspaceSlug) { } /** - * Returns workspace members and optionally excludes members whose usernames are in excludeUsernames. + * Returns workspace members and optionally excludes members whose usernames are + * in excludeUsernames. * excludeUsernames can be null or empty (no exclusions). */ @Transactional(readOnly = true) @@ -209,18 +215,19 @@ public EWorkspaceRole getUserRole(String workspaceSlug, Long userId) { @Transactional public void deleteWorkspace(Long actorUserId, String workspaceSlug) { Workspace workspace = getWorkspaceBySlug(workspaceSlug); - + // Verify actor is the owner - WorkspaceMember actorMember = workspaceMemberRepository.findByWorkspaceIdAndUserId(workspace.getId(), actorUserId) + WorkspaceMember actorMember = workspaceMemberRepository + .findByWorkspaceIdAndUserId(workspace.getId(), actorUserId) .orElseThrow(() -> new SecurityException("User is not a member of this workspace")); - + if (actorMember.getRole() != EWorkspaceRole.OWNER) { throw new SecurityException("Only the workspace owner can delete a workspace"); } - + // Delete all workspace members first workspaceMemberRepository.deleteAllByWorkspaceId(workspace.getId()); - + // Delete the workspace workspaceRepository.delete(workspace); } @@ -231,19 +238,20 @@ public void deleteWorkspace(Long actorUserId, String workspaceSlug) { @Transactional public Workspace scheduleDeletion(Long actorUserId, String workspaceSlug) { Workspace workspace = getWorkspaceBySlug(workspaceSlug); - + // Verify actor is the owner - WorkspaceMember actorMember = workspaceMemberRepository.findByWorkspaceIdAndUserId(workspace.getId(), actorUserId) + WorkspaceMember actorMember = workspaceMemberRepository + .findByWorkspaceIdAndUserId(workspace.getId(), actorUserId) .orElseThrow(() -> new SecurityException("User is not a member of this workspace")); - + if (actorMember.getRole() != EWorkspaceRole.OWNER) { throw new SecurityException("Only the workspace owner can schedule workspace deletion"); } - + if (workspace.isScheduledForDeletion()) { throw new IllegalStateException("Workspace is already scheduled for deletion"); } - + workspace.scheduleDeletion(actorUserId); return workspaceRepository.save(workspace); } @@ -254,19 +262,20 @@ public Workspace scheduleDeletion(Long actorUserId, String workspaceSlug) { @Transactional public Workspace cancelScheduledDeletion(Long actorUserId, String workspaceSlug) { Workspace workspace = getWorkspaceBySlug(workspaceSlug); - + // Verify actor is the owner - WorkspaceMember actorMember = workspaceMemberRepository.findByWorkspaceIdAndUserId(workspace.getId(), actorUserId) + WorkspaceMember actorMember = workspaceMemberRepository + .findByWorkspaceIdAndUserId(workspace.getId(), actorUserId) .orElseThrow(() -> new SecurityException("User is not a member of this workspace")); - + if (actorMember.getRole() != EWorkspaceRole.OWNER) { throw new SecurityException("Only the workspace owner can cancel workspace deletion"); } - + if (!workspace.isScheduledForDeletion()) { throw new IllegalStateException("Workspace is not scheduled for deletion"); } - + workspace.cancelDeletion(); return workspaceRepository.save(workspace); } diff --git a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/ai/AIConnectionCrudIT.java b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/ai/AIConnectionCrudIT.java index 2c94d528..69d3bef6 100644 --- a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/ai/AIConnectionCrudIT.java +++ b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/ai/AIConnectionCrudIT.java @@ -36,8 +36,7 @@ void shouldCreateOpenRouterConnection() { { "providerKey": "OPENROUTER", "aiModel": "anthropic/claude-3-haiku", - "apiKey": "test-api-key-openrouter", - "tokenLimitation": "200000" + "apiKey": "test-api-key-openrouter" } """; @@ -50,7 +49,6 @@ void shouldCreateOpenRouterConnection() { .statusCode(201) .body("providerKey", equalTo("OPENROUTER")) .body("aiModel", equalTo("anthropic/claude-3-haiku")) - .body("tokenLimitation", equalTo(200000)) .body("id", notNullValue()); } @@ -62,8 +60,7 @@ void shouldCreateOpenAIConnection() { { "providerKey": "OPENAI", "aiModel": "gpt-4o-mini", - "apiKey": "test-api-key-openai", - "tokenLimitation": "128000" + "apiKey": "test-api-key-openai" } """; @@ -86,8 +83,7 @@ void shouldCreateAnthropicConnection() { { "providerKey": "ANTHROPIC", "aiModel": "claude-3-haiku-20240307", - "apiKey": "test-api-key-anthropic", - "tokenLimitation": "200000" + "apiKey": "test-api-key-anthropic" } """; @@ -110,8 +106,7 @@ void shouldListAIConnections() { { "providerKey": "OPENROUTER", "aiModel": "test-model", - "apiKey": "test-key", - "tokenLimitation": "100000" + "apiKey": "test-key" } """; @@ -140,8 +135,7 @@ void shouldUpdateAIConnection() { { "providerKey": "OPENROUTER", "aiModel": "original-model", - "apiKey": "original-key", - "tokenLimitation": "100000" + "apiKey": "original-key" } """; @@ -158,8 +152,7 @@ void shouldUpdateAIConnection() { { "providerKey": "OPENROUTER", "aiModel": "updated-model", - "apiKey": "updated-key", - "tokenLimitation": "150000" + "apiKey": "updated-key" } """; @@ -170,8 +163,7 @@ void shouldUpdateAIConnection() { .patch("/api/{workspaceSlug}/ai/{connectionId}", testWorkspace.getSlug(), connectionId) .then() .statusCode(200) - .body("aiModel", equalTo("updated-model")) - .body("tokenLimitation", equalTo(150000)); + .body("aiModel", equalTo("updated-model")); } @Test @@ -182,8 +174,7 @@ void shouldDeleteAIConnection() { { "providerKey": "OPENROUTER", "aiModel": "to-delete", - "apiKey": "delete-key", - "tokenLimitation": "100000" + "apiKey": "delete-key" } """; @@ -212,8 +203,7 @@ void shouldRequireAdminRightsForAIOperations() { { "providerKey": "OPENROUTER", "aiModel": "test-model", - "apiKey": "test-key", - "tokenLimitation": "100000" + "apiKey": "test-key" } """; @@ -272,8 +262,7 @@ void shouldValidateProviderKey() { { "providerKey": "INVALID_PROVIDER", "aiModel": "test-model", - "apiKey": "test-key", - "tokenLimitation": "100000" + "apiKey": "test-key" } """; @@ -300,8 +289,7 @@ void shouldPreventCrossWorkspaceAccess() { { "providerKey": "OPENROUTER", "aiModel": "other-ws-model", - "apiKey": "other-ws-key", - "tokenLimitation": "100000" + "apiKey": "other-ws-key" } """; diff --git a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/auth/UserAuthFlowIT.java b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/auth/UserAuthFlowIT.java index 5d0ffaca..bb8d4220 100644 --- a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/auth/UserAuthFlowIT.java +++ b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/auth/UserAuthFlowIT.java @@ -313,8 +313,7 @@ void shouldHandleWorkspaceRoleDowngrade() { { "providerKey": "OPENROUTER", "aiModel": "test", - "apiKey": "key", - "tokenLimitation": "100000" + "apiKey": "key" } """; given() diff --git a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/builder/AIConnectionBuilder.java b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/builder/AIConnectionBuilder.java index ef858489..69ecb04e 100644 --- a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/builder/AIConnectionBuilder.java +++ b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/builder/AIConnectionBuilder.java @@ -14,7 +14,6 @@ public class AIConnectionBuilder { private AIProviderKey providerKey = AIProviderKey.OPENROUTER; private String aiModel = "anthropic/claude-3-haiku"; private String apiKeyEncrypted = "encrypted-test-key"; - private int tokenLimitation = 200000; public static AIConnectionBuilder anAIConnection() { return new AIConnectionBuilder(); @@ -45,11 +44,6 @@ public AIConnectionBuilder withApiKeyEncrypted(String apiKeyEncrypted) { return this; } - public AIConnectionBuilder withTokenLimitation(int tokenLimitation) { - this.tokenLimitation = tokenLimitation; - return this; - } - public AIConnectionBuilder openAI() { this.providerKey = AIProviderKey.OPENAI; this.aiModel = "gpt-4o-mini"; @@ -75,7 +69,6 @@ public AIConnection build() { connection.setProviderKey(providerKey); connection.setAiModel(aiModel); connection.setApiKeyEncrypted(apiKeyEncrypted); - connection.setTokenLimitation(tokenLimitation); return connection; } } diff --git a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/util/AuthTestHelper.java b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/util/AuthTestHelper.java index c819e219..8902eaf5 100644 --- a/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/util/AuthTestHelper.java +++ b/java-ecosystem/tests/integration-tests/src/test/java/org/rostilos/codecrow/integration/util/AuthTestHelper.java @@ -254,7 +254,6 @@ public AIConnection createTestAiConnection(Workspace workspace, String name, AIP aiConnection.setProviderKey(provider); aiConnection.setAiModel("gpt-4"); aiConnection.setApiKeyEncrypted("test-encrypted-api-key-" + UUID.randomUUID().toString().substring(0, 8)); - aiConnection.setTokenLimitation(100000); return aiConnectionRepository.save(aiConnection); } diff --git a/python-ecosystem/mcp-client/Dockerfile b/python-ecosystem/mcp-client/Dockerfile index 9dd14dba..2c03cc27 100644 --- a/python-ecosystem/mcp-client/Dockerfile +++ b/python-ecosystem/mcp-client/Dockerfile @@ -22,6 +22,7 @@ RUN pip install --no-cache-dir -r requirements.txt # --- Builder Stage 2: Copy and Install Application Modules --- # Copy application source code and the JAR COPY main.py . +COPY api ./api/ COPY server ./server/ COPY model ./model/ COPY service ./service/ @@ -56,6 +57,7 @@ COPY --from=builder /usr/local/lib/python3.11/site-packages /usr/local/lib/pytho # Copy application code from the builder stage COPY --from=builder /app/main.py ./ +COPY --from=builder /app/api ./api/ COPY --from=builder /app/server ./server/ COPY --from=builder /app/model ./model/ COPY --from=builder /app/service ./service/ diff --git a/python-ecosystem/mcp-client/api/__init__.py b/python-ecosystem/mcp-client/api/__init__.py new file mode 100644 index 00000000..613c59d6 --- /dev/null +++ b/python-ecosystem/mcp-client/api/__init__.py @@ -0,0 +1,8 @@ +""" +API Package. + +Contains the FastAPI application and routers. +""" +from api.app import create_app, run_http_server + +__all__ = ["create_app", "run_http_server"] diff --git a/python-ecosystem/mcp-client/api/app.py b/python-ecosystem/mcp-client/api/app.py new file mode 100644 index 00000000..fc513bc4 --- /dev/null +++ b/python-ecosystem/mcp-client/api/app.py @@ -0,0 +1,73 @@ +""" +FastAPI Application Factory. + +Creates and configures the FastAPI application with all routers. +Uses lifespan context manager for proper startup/shutdown of shared resources. +""" +import os +import logging +from contextlib import asynccontextmanager +from fastapi import FastAPI + +from api.routers import health, review, commands +from api.middleware import ServiceSecretMiddleware +from service.review.review_service import ReviewService +from service.command.command_service import CommandService + +logger = logging.getLogger(__name__) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Manage application lifecycle: create services on startup, clean up on shutdown.""" + # --- Startup --- + logger.info("Initializing application services...") + review_service = ReviewService() + command_service = CommandService() + + app.state.review_service = review_service + app.state.command_service = command_service + logger.info("Application services ready") + + yield + + # --- Shutdown --- + logger.info("Shutting down application services...") + # Close the RagClient HTTP pools owned by each service + try: + await review_service.rag_client.close() + except Exception as e: + logger.warning(f"Error closing review RagClient: {e}") + try: + await command_service.rag_client.close() + except Exception as e: + logger.warning(f"Error closing command RagClient: {e}") + logger.info("Application services shut down") + + +def create_app() -> FastAPI: + """Create and configure FastAPI application.""" + app = FastAPI(title="codecrow-mcp-client", lifespan=lifespan) + + # Service-to-service auth + app.add_middleware(ServiceSecretMiddleware) + + # Register routers + app.include_router(health.router) + app.include_router(review.router) + app.include_router(commands.router) + + return app + + +def run_http_server(host: str = "0.0.0.0", port: int = 8000): + """Run the FastAPI application.""" + app = create_app() + import uvicorn + uvicorn.run(app, host=host, port=port, log_level="info", timeout_keep_alive=300) + + +if __name__ == "__main__": + host = os.environ.get("AI_CLIENT_HOST", "0.0.0.0") + port = int(os.environ.get("AI_CLIENT_PORT", "8000")) + run_http_server(host=host, port=port) diff --git a/python-ecosystem/mcp-client/api/middleware.py b/python-ecosystem/mcp-client/api/middleware.py new file mode 100644 index 00000000..96bef76c --- /dev/null +++ b/python-ecosystem/mcp-client/api/middleware.py @@ -0,0 +1,56 @@ +""" +Shared-secret authentication middleware. + +Validates that internal service-to-service requests carry +the correct X-Service-Secret header matching the SERVICE_SECRET env var. +""" +import os +import logging + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse + +logger = logging.getLogger(__name__) + +# Paths that skip auth (health checks, readiness probes) +_PUBLIC_PATHS = frozenset({"/health", "/docs", "/openapi.json", "/redoc"}) + + +class ServiceSecretMiddleware(BaseHTTPMiddleware): + """Reject requests that don't carry a valid shared service secret.""" + + def __init__(self, app, secret: str | None = None): + super().__init__(app) + self.secret = secret or os.environ.get("SERVICE_SECRET", "") + if self.secret: + logger.info("ServiceSecretMiddleware: secret configured (length=%d)", len(self.secret)) + else: + logger.warning("ServiceSecretMiddleware: no secret configured — auth disabled") + + async def dispatch(self, request: Request, call_next): + # Skip auth for health/doc endpoints + if request.url.path in _PUBLIC_PATHS: + return await call_next(request) + + # If no secret is configured, allow all (dev mode) + if not self.secret: + return await call_next(request) + + provided = request.headers.get("x-service-secret", "") + if provided != self.secret: + logger.warning( + "Unauthorized request to %s from %s — " + "provided_len=%d expected_len=%d match=%s", + request.url.path, + request.client.host if request.client else "unknown", + len(provided), + len(self.secret), + provided == self.secret, + ) + return JSONResponse( + status_code=401, + content={"detail": "Missing or invalid service secret"}, + ) + + return await call_next(request) diff --git a/python-ecosystem/mcp-client/api/routers/__init__.py b/python-ecosystem/mcp-client/api/routers/__init__.py new file mode 100644 index 00000000..06f12462 --- /dev/null +++ b/python-ecosystem/mcp-client/api/routers/__init__.py @@ -0,0 +1,5 @@ +""" +API Routers Package. + +Contains FastAPI routers for different endpoint groups. +""" diff --git a/python-ecosystem/mcp-client/api/routers/commands.py b/python-ecosystem/mcp-client/api/routers/commands.py new file mode 100644 index 00000000..de6046e6 --- /dev/null +++ b/python-ecosystem/mcp-client/api/routers/commands.py @@ -0,0 +1,198 @@ +""" +Command API endpoints (summarize, ask). +""" +import json +import asyncio +from typing import Dict, Any +from fastapi import APIRouter, Request +from starlette.responses import StreamingResponse + +from model.dtos import ( + SummarizeRequestDto, SummarizeResponseDto, + AskRequestDto, AskResponseDto, +) +from service.command.command_service import CommandService + +router = APIRouter(tags=["commands"]) + + +def get_command_service(request: Request) -> CommandService: + """Retrieve the CommandService instance created during app lifespan.""" + return request.app.state.command_service + + +@router.post("/review/summarize", response_model=SummarizeResponseDto) +async def summarize_endpoint(req: SummarizeRequestDto, request: Request): + """ + HTTP endpoint to process /codecrow summarize command. + + Generates a comprehensive PR summary with: + - Overview of changes + - Key files modified + - Impact analysis + - Architecture diagram (Mermaid or ASCII) + """ + command_service = get_command_service(request) + + try: + wants_stream = _wants_streaming(request) + + if not wants_stream: + # Non-streaming behavior + result = await command_service.process_summarize(req) + return SummarizeResponseDto( + summary=result.get("summary"), + diagram=result.get("diagram"), + diagramType=result.get("diagramType", "MERMAID"), + error=result.get("error") + ) + + # Streaming behavior + async def event_stream(): + queue = asyncio.Queue() + + yield _json_event({"type": "status", "state": "queued", "message": "summarize request received"}) + + def event_callback(event: Dict[str, Any]): + try: + queue.put_nowait(event) + except asyncio.QueueFull: + pass + + async def runner(): + try: + result = await command_service.process_summarize(req, event_callback=event_callback) + await queue.put({ + "type": "final", + "result": result + }) + except Exception as e: + await queue.put({"type": "error", "message": str(e)}) + + task = asyncio.create_task(runner()) + + async for event in _drain_queue_until_final(queue, task): + yield _json_event(event) + + return StreamingResponse(event_stream(), media_type="application/x-ndjson") + + except Exception as e: + return SummarizeResponseDto(error=f"Summarize failed: {str(e)}") + + +@router.post("/review/ask", response_model=AskResponseDto) +async def ask_endpoint(req: AskRequestDto, request: Request): + """ + HTTP endpoint to process /codecrow ask command. + + Answers questions about: + - Specific issues + - PR changes + - Codebase (using RAG) + - Analysis results + """ + command_service = get_command_service(request) + + try: + wants_stream = _wants_streaming(request) + + if not wants_stream: + # Non-streaming behavior + result = await command_service.process_ask(req) + return AskResponseDto( + answer=result.get("answer"), + error=result.get("error") + ) + + # Streaming behavior + async def event_stream(): + queue = asyncio.Queue() + + yield _json_event({"type": "status", "state": "queued", "message": "ask request received"}) + + def event_callback(event: Dict[str, Any]): + try: + queue.put_nowait(event) + except asyncio.QueueFull: + pass + + async def runner(): + try: + result = await command_service.process_ask(req, event_callback=event_callback) + await queue.put({ + "type": "final", + "result": result + }) + except Exception as e: + await queue.put({"type": "error", "message": str(e)}) + + task = asyncio.create_task(runner()) + + async for event in _drain_queue_until_final(queue, task): + yield _json_event(event) + + return StreamingResponse(event_stream(), media_type="application/x-ndjson") + + except Exception as e: + return AskResponseDto(error=f"Ask failed: {str(e)}") + + +def _wants_streaming(request: Request) -> bool: + """Check if client wants streaming response.""" + accept_header = request.headers.get("accept", "") + return "application/x-ndjson" in accept_header.lower() + + +def _json_event(event: Dict[str, Any]) -> str: + """Serialize event to NDJSON line.""" + return json.dumps(event) + "\n" + + +async def _drain_queue_until_final(queue: asyncio.Queue, task: asyncio.Task): + """ + Drain events from queue until we see a final/error event and task is done. + Yields each event as it arrives. + """ + seen_final = False + + while True: + try: + # Wait for event with timeout to prevent hanging + event = await asyncio.wait_for(queue.get(), timeout=1.0) + yield event + + # Check if this is a terminal event + event_type = event.get("type") + if event_type in ("final", "error"): + seen_final = True + # Continue draining in case there are more events + + # If we've seen final and task is done, check for remaining events then exit + if seen_final and task.done(): + # Give a moment for any last events + await asyncio.sleep(0.1) + try: + while True: + event = queue.get_nowait() + yield event + except asyncio.QueueEmpty: + break + break + + except asyncio.TimeoutError: + # No event available, check if task is done + if task.done(): + # Task finished, drain any remaining events + try: + while True: + event = queue.get_nowait() + yield event + if event.get("type") in ("final", "error"): + seen_final = True + except asyncio.QueueEmpty: + pass + + # If we saw a final event or no more events, we're done + if seen_final or task.done(): + break + # Otherwise continue waiting for events diff --git a/python-ecosystem/mcp-client/api/routers/health.py b/python-ecosystem/mcp-client/api/routers/health.py new file mode 100644 index 00000000..98c53522 --- /dev/null +++ b/python-ecosystem/mcp-client/api/routers/health.py @@ -0,0 +1,12 @@ +""" +Health check endpoints. +""" +from fastapi import APIRouter + +router = APIRouter(tags=["health"]) + + +@router.get("/health") +def health(): + """Health check endpoint.""" + return {"status": "ok"} diff --git a/python-ecosystem/mcp-client/api/routers/review.py b/python-ecosystem/mcp-client/api/routers/review.py new file mode 100644 index 00000000..ca167662 --- /dev/null +++ b/python-ecosystem/mcp-client/api/routers/review.py @@ -0,0 +1,151 @@ +""" +Review API endpoints. +""" +import json +import asyncio +from typing import Dict, Any +from fastapi import APIRouter, Request +from starlette.responses import StreamingResponse + +from model.dtos import ReviewRequestDto, ReviewResponseDto +from service.review.review_service import ReviewService +from utils.response_parser import ResponseParser + +router = APIRouter(tags=["review"]) + + +def get_review_service(request: Request) -> ReviewService: + """Retrieve the ReviewService instance created during app lifespan.""" + return request.app.state.review_service + + +@router.post("/review", response_model=ReviewResponseDto) +async def review_endpoint(req: ReviewRequestDto, request: Request): + """ + HTTP endpoint to accept review requests from the pipeline agent. + + Behavior: + - If the client requests streaming via header `Accept: application/x-ndjson`, + the endpoint will return a StreamingResponse that yields NDJSON events as they occur. + - Otherwise it preserves the original behavior and returns a single ReviewResponseDto JSON body. + """ + review_service = get_review_service(request) + + try: + wants_stream = _wants_streaming(request) + + if not wants_stream: + # Non-streaming (legacy) behavior + result = await review_service.process_review_request(req) + return ReviewResponseDto( + result=result.get("result"), + error=result.get("error") + ) + + # Streaming behavior + async def event_stream(): + queue = asyncio.Queue() + + # Emit initial queued status + yield _json_event({"type": "status", "state": "queued", "message": "request received"}) + + # Event callback to capture service events + def event_callback(event: Dict[str, Any]): + try: + queue.put_nowait(event) + except asyncio.QueueFull: + pass # Skip if queue is full + + # Run processing in background + async def runner(): + try: + result = await review_service.process_review_request( + req, + event_callback=event_callback + ) + # Emit final event with result + final_event = { + "type": "final", + "result": result.get("result") + } + await queue.put(final_event) + except Exception as e: + await queue.put({ + "type": "error", + "message": str(e) + }) + + task = asyncio.create_task(runner()) + + # Drain queue and yield events + async for event in _drain_queue_until_final(queue, task): + yield _json_event(event) + + return StreamingResponse(event_stream(), media_type="application/x-ndjson") + + except Exception as e: + error_response = ResponseParser.create_error_response( + "HTTP request processing failed", str(e) + ) + return ReviewResponseDto(result=error_response) + + +def _wants_streaming(request: Request) -> bool: + """Check if client wants streaming response.""" + accept_header = request.headers.get("accept", "") + return "application/x-ndjson" in accept_header.lower() + + +def _json_event(event: Dict[str, Any]) -> str: + """Serialize event to NDJSON line.""" + return json.dumps(event) + "\n" + + +async def _drain_queue_until_final(queue: asyncio.Queue, task: asyncio.Task): + """ + Drain events from queue until we see a final/error event and task is done. + Yields each event as it arrives. + """ + seen_final = False + + while True: + try: + # Wait for event with timeout to prevent hanging + event = await asyncio.wait_for(queue.get(), timeout=1.0) + yield event + + # Check if this is a terminal event + event_type = event.get("type") + if event_type in ("final", "error"): + seen_final = True + # Continue draining in case there are more events + + # If we've seen final and task is done, check for remaining events then exit + if seen_final and task.done(): + # Give a moment for any last events + await asyncio.sleep(0.1) + try: + while True: + event = queue.get_nowait() + yield event + except asyncio.QueueEmpty: + break + break + + except asyncio.TimeoutError: + # No event available, check if task is done + if task.done(): + # Task finished, drain any remaining events + try: + while True: + event = queue.get_nowait() + yield event + if event.get("type") in ("final", "error"): + seen_final = True + except asyncio.QueueEmpty: + pass + + # If we saw a final event or no more events, we're done + if seen_final or task.done(): + break + # Otherwise continue waiting for events diff --git a/python-ecosystem/mcp-client/llm/llm_factory.py b/python-ecosystem/mcp-client/llm/llm_factory.py index 470e7c4e..2b199fa6 100644 --- a/python-ecosystem/mcp-client/llm/llm_factory.py +++ b/python-ecosystem/mcp-client/llm/llm_factory.py @@ -139,6 +139,9 @@ def create_llm(ai_model: str, ai_provider: str, ai_api_key: str, temperature: Op # Normalize provider provider = LLMFactory._normalize_provider(ai_provider) + # CRITICAL: Log the model being used for debugging + logger.info(f"Creating LLM instance: provider={provider}, model={ai_model}, temperature={temperature}") + # Check for unsupported Gemini thinking models (applies to all providers) LLMFactory._check_unsupported_gemini_model(ai_model) @@ -196,8 +199,7 @@ def create_llm(ai_model: str, ai_provider: str, ai_api_key: str, temperature: Op return ChatGoogleGenerativeAI( google_api_key=ai_api_key, model=ai_model, - temperature=temperature if temperature > 0 else 0.1, - thinking_level="minimal", + temperature=temperature if temperature > 0 else 0.1 ) else: # Gemini 2.x models use thinking_budget parameter diff --git a/python-ecosystem/mcp-client/main.py b/python-ecosystem/mcp-client/main.py index abb28d6f..d62dc50b 100644 --- a/python-ecosystem/mcp-client/main.py +++ b/python-ecosystem/mcp-client/main.py @@ -15,9 +15,9 @@ import sys import logging import warnings -import threading from server.stdin_handler import StdinHandler +from api.app import run_http_server # Configure logging - only if not already configured # This prevents duplicate handlers when libraries also configure logging @@ -38,38 +38,24 @@ # Suppress pydantic warnings warnings.filterwarnings('ignore', category=UserWarning, module='pydantic') -# Wrap stderr to filter out JSONRPC parsing errors from MCP library -# These occur when Java MCP servers leak log messages to stdout -class FilteredStderr: - def __init__(self, original_stderr): - self.original_stderr = original_stderr - self.buffer = "" - self._lock = threading.Lock() - self._suppress_next_lines = 0 - def write(self, text): - with self._lock: - # Check if this is the start of a JSONRPC parsing error - if "Failed to parse JSONRPC message from server" in text: - self._suppress_next_lines = 15 # Suppress the next ~15 lines (traceback) - return +class _McpJsonRpcFilter(logging.Filter): + """Filter out noisy JSONRPC parsing errors from the mcp_use library. + + These occur when Java MCP servers leak log messages to stdout, which + the mcp_use library then fails to parse as JSON-RPC. They are harmless + and clutter the logs. + """ + def filter(self, record: logging.LogRecord) -> bool: + msg = record.getMessage() + if "Failed to parse JSONRPC message from server" in msg: + return False + return True - # If we're suppressing lines, decrement counter - if self._suppress_next_lines > 0: - self._suppress_next_lines -= 1 - return - # Otherwise, write to original stderr - self.original_stderr.write(text) - - def flush(self): - self.original_stderr.flush() - - def __getattr__(self, name): - return getattr(self.original_stderr, name) - -# Install the filtered stderr wrapper -sys.stderr = FilteredStderr(sys.stderr) +# Install the filter on the mcp_use loggers instead of wrapping stderr +for _logger_name in ("mcp_use", "mcp_use.client", "mcp_use.client.session"): + logging.getLogger(_logger_name).addFilter(_McpJsonRpcFilter()) def main(): @@ -79,7 +65,6 @@ def main(): handler.process_stdin_request() else: # Run in HTTP server mode (default) - from server.web_server import run_http_server host = os.environ.get("AI_CLIENT_HOST", "0.0.0.0") port = int(os.environ.get("AI_CLIENT_PORT", "8000")) run_http_server(host=host, port=port) diff --git a/python-ecosystem/mcp-client/model/__init__.py b/python-ecosystem/mcp-client/model/__init__.py new file mode 100644 index 00000000..5e9aa368 --- /dev/null +++ b/python-ecosystem/mcp-client/model/__init__.py @@ -0,0 +1,94 @@ +""" +Model package - Re-exports all models for backward compatibility. + +The models are split into logical modules: +- enums: IssueCategory, AnalysisMode, RelationshipType +- enrichment: File enrichment DTOs (FileContentDto, PrEnrichmentDataDto, etc.) +- dtos: Request/Response DTOs (ReviewRequestDto, SummarizeRequestDto, etc.) +- output_schemas: MCP Agent output schemas (CodeReviewOutput, CodeReviewIssue, etc.) +- multi_stage: Multi-stage review models (ReviewPlan, FileReviewOutput, etc.) +""" + +# Enums +from model.enums import ( + IssueCategory, + AnalysisMode, + RelationshipType, +) + +# Enrichment models +from model.enrichment import ( + FileContentDto, + ParsedFileMetadataDto, + FileRelationshipDto, + EnrichmentStats, + PrEnrichmentDataDto, +) + +# DTOs +from model.dtos import ( + IssueDTO, + ReviewRequestDto, + ReviewResponseDto, + SummarizeRequestDto, + SummarizeResponseDto, + AskRequestDto, + AskResponseDto, +) + +# Output schemas +from model.output_schemas import ( + CodeReviewIssue, + CodeReviewOutput, + SummarizeOutput, + AskOutput, +) + +# Multi-stage review models +from model.multi_stage import ( + FileReviewOutput, + FileReviewBatchOutput, + ReviewFile, + FileGroup, + FileToSkip, + ReviewPlan, + CrossFileIssue, + DataFlowConcern, + CrossFileAnalysisResult, +) + +__all__ = [ + # Enums + "IssueCategory", + "AnalysisMode", + "RelationshipType", + # Enrichment + "FileContentDto", + "ParsedFileMetadataDto", + "FileRelationshipDto", + "EnrichmentStats", + "PrEnrichmentDataDto", + # DTOs + "IssueDTO", + "ReviewRequestDto", + "ReviewResponseDto", + "SummarizeRequestDto", + "SummarizeResponseDto", + "AskRequestDto", + "AskResponseDto", + # Output schemas + "CodeReviewIssue", + "CodeReviewOutput", + "SummarizeOutput", + "AskOutput", + # Multi-stage + "FileReviewOutput", + "FileReviewBatchOutput", + "ReviewFile", + "FileGroup", + "FileToSkip", + "ReviewPlan", + "CrossFileIssue", + "DataFlowConcern", + "CrossFileAnalysisResult", +] diff --git a/python-ecosystem/mcp-client/model/models.py b/python-ecosystem/mcp-client/model/dtos.py similarity index 51% rename from python-ecosystem/mcp-client/model/models.py rename to python-ecosystem/mcp-client/model/dtos.py index 153de871..192369fc 100644 --- a/python-ecosystem/mcp-client/model/models.py +++ b/python-ecosystem/mcp-client/model/dtos.py @@ -1,27 +1,8 @@ -from typing import Optional, Any, Dict, List +from typing import Optional, Any, List from pydantic import BaseModel, Field, AliasChoices from datetime import datetime -from enum import Enum - -class IssueCategory(str, Enum): - """Valid issue categories for code analysis.""" - SECURITY = "SECURITY" - PERFORMANCE = "PERFORMANCE" - CODE_QUALITY = "CODE_QUALITY" - BUG_RISK = "BUG_RISK" - STYLE = "STYLE" - DOCUMENTATION = "DOCUMENTATION" - BEST_PRACTICES = "BEST_PRACTICES" - ERROR_HANDLING = "ERROR_HANDLING" - TESTING = "TESTING" - ARCHITECTURE = "ARCHITECTURE" - - -class AnalysisMode(str, Enum): - """Analysis mode for PR reviews.""" - FULL = "FULL" # Full PR diff analysis (first review or escalation) - INCREMENTAL = "INCREMENTAL" # Delta diff analysis (subsequent reviews) +from model.enrichment import PrEnrichmentDataDto class IssueDTO(BaseModel): @@ -41,6 +22,11 @@ class IssueDTO(BaseModel): branch: Optional[str] = None pullRequestId: Optional[str] = None status: Optional[str] = None # open|resolved|ignored + # Resolution tracking fields (for full PR issue history) + prVersion: Optional[int] = None # Which PR iteration this issue was found in + resolvedDescription: Optional[str] = None # How the issue was resolved + resolvedByCommit: Optional[str] = None # Commit hash that resolved the issue + resolvedInPrVersion: Optional[int] = None # PR version where this was resolved # Legacy fields for backwards compatibility title: Optional[str] = None # Legacy - use reason instead description: Optional[str] = None # Legacy - use suggestedFixDescription instead @@ -84,12 +70,16 @@ class ReviewRequestDto(BaseModel): deltaDiff: Optional[str] = Field(default=None, description="Delta diff between previous and current commit (only for INCREMENTAL mode)") previousCommitHash: Optional[str] = Field(default=None, description="Previously analyzed commit hash") currentCommitHash: Optional[str] = Field(default=None, description="Current commit hash being analyzed") + # File enrichment data (full file contents + pre-computed dependency graph) + enrichmentData: Optional[PrEnrichmentDataDto] = Field(default=None, description="Pre-computed file contents and dependency relationships from Java") + class ReviewResponseDto(BaseModel): result: Optional[Any] = None error: Optional[str] = None exception: Optional[str] = None + class SummarizeRequestDto(BaseModel): """Request model for PR summarization command.""" projectId: int @@ -111,6 +101,7 @@ class SummarizeRequestDto(BaseModel): maxAllowedTokens: Optional[int] = None vcsProvider: Optional[str] = Field(default=None, description="VCS provider type (github, bitbucket_cloud)") + class SummarizeResponseDto(BaseModel): """Response model for PR summarization command.""" summary: Optional[str] = None @@ -118,6 +109,7 @@ class SummarizeResponseDto(BaseModel): diagramType: Optional[str] = Field(default="MERMAID", description="MERMAID or ASCII") error: Optional[str] = None + class AskRequestDto(BaseModel): """Request model for ask command.""" projectId: int @@ -140,138 +132,8 @@ class AskRequestDto(BaseModel): analysisContext: Optional[str] = Field(default=None, description="Existing analysis data for context") issueReferences: Optional[List[str]] = Field(default_factory=list, description="Issue IDs referenced in the question") + class AskResponseDto(BaseModel): """Response model for ask command.""" answer: Optional[str] = None error: Optional[str] = None - -# ==================== Output Schemas for MCP Agent ==================== -# These Pydantic models are used with MCPAgent's output_schema parameter -# to ensure structured JSON output from the LLM. - -class CodeReviewIssue(BaseModel): - """Schema for a single code review issue.""" - # Optional issue identifier (preserve DB/client-side ids for reconciliation) - id: Optional[str] = Field(default=None, description="Optional issue id to link to existing issues") - severity: str = Field(description="Issue severity: HIGH, MEDIUM, LOW, or INFO") - category: str = Field(description="Issue category: SECURITY, PERFORMANCE, CODE_QUALITY, BUG_RISK, STYLE, DOCUMENTATION, BEST_PRACTICES, ERROR_HANDLING, TESTING, or ARCHITECTURE") - file: str = Field(description="File path where the issue is located") - line: str = Field(description="Line number or range (e.g., '42' or '42-45')") - reason: str = Field(description="Clear explanation of the issue") - suggestedFixDescription: str = Field(description="Description of the suggested fix") - suggestedFixDiff: Optional[str] = Field(default=None, description="Optional unified diff format patch for the fix") - isResolved: bool = Field(default=False, description="Whether this issue from previous analysis is resolved") - # Additional fields preserved from previous issues during reconciliation - visibility: Optional[str] = Field(default=None, description="Issue visibility status") - codeSnippet: Optional[str] = Field(default=None, description="Code snippet associated with the issue") - - -class CodeReviewOutput(BaseModel): - """Schema for the complete code review output.""" - comment: str = Field(description="High-level summary of the PR analysis with key findings and recommendations") - issues: List[CodeReviewIssue] = Field(default_factory=list, description="List of identified issues in the code") - - -class SummarizeOutput(BaseModel): - """Schema for PR summarization output.""" - summary: str = Field(description="Comprehensive summary of the PR changes, purpose, and impact") - diagram: str = Field(default="", description="Visual diagram of the changes (Mermaid or ASCII format)") - diagramType: str = Field(default="MERMAID", description="Type of diagram: MERMAID or ASCII") - - -class AskOutput(BaseModel): - """Schema for ask command output.""" - answer: str = Field(description="Well-formatted markdown answer to the user's question") - - -# ==================== Multi-Stage Review Models ==================== - -class FileReviewOutput(BaseModel): - """Stage 1 Output: Single file review result.""" - file: str - analysis_summary: str - issues: List[CodeReviewIssue] = Field(default_factory=list) - confidence: str = Field(description="Confidence level (HIGH/MEDIUM/LOW)") - note: str = Field(default="", description="Optional analysis note") - - - -class FileReviewBatchOutput(BaseModel): - """Stage 1 Output: Batch of file reviews.""" - reviews: List[FileReviewOutput] = Field(description="List of review results for the files in the batch") - - - -class ReviewFile(BaseModel): - """File details for review planning.""" - path: str - focus_areas: List[str] = Field(default_factory=list, description="Specific areas to focus on (SECURITY, ARCHITECTURE, etc.)") - risk_level: str = Field(description="CRITICAL, HIGH, MEDIUM, or LOW") - estimated_issues: Optional[int] = Field(default=0) - - -class FileGroup(BaseModel): - """Group of files to be reviewed together.""" - group_id: str - priority: str = Field(description="CRITICAL, HIGH, MEDIUM, LOW") - rationale: str - files: List[ReviewFile] - - -class FileToSkip(BaseModel): - """File skipped from deep review.""" - path: str - reason: str - - -class ReviewPlan(BaseModel): - """Stage 0 Output: Plan for the review scanning.""" - analysis_summary: str - file_groups: List[FileGroup] - files_to_skip: List[FileToSkip] = Field(default_factory=list) - cross_file_concerns: List[str] = Field(default_factory=list, description="Hypotheses to verify in Stage 2") - - -class CrossFileIssue(BaseModel): - """Issue spanning multiple files (Stage 2).""" - id: str - severity: str - category: str - title: str - affected_files: List[str] - description: str - evidence: str - business_impact: str - suggestion: str - - -class DataFlowConcern(BaseModel): - """Stage 2: Data flow gap analysis.""" - flow: str - gap: str - files_involved: List[str] - severity: str - - -class ImmutabilityCheck(BaseModel): - """Stage 2: Immutability usage check.""" - rule: str - check_pass: bool = Field(alias="check_pass") - evidence: str - - -class DatabaseIntegrityCheck(BaseModel): - """Stage 2: DB integrity check.""" - concerns: List[str] - findings: List[str] - - -class CrossFileAnalysisResult(BaseModel): - """Stage 2 Output: Cross-file architectural analysis.""" - pr_risk_level: str - cross_file_issues: List[CrossFileIssue] - data_flow_concerns: List[DataFlowConcern] = Field(default_factory=list) - immutability_enforcement: Optional[ImmutabilityCheck] = None - database_integrity: Optional[DatabaseIntegrityCheck] = None - pr_recommendation: str - confidence: str \ No newline at end of file diff --git a/python-ecosystem/mcp-client/model/enrichment.py b/python-ecosystem/mcp-client/model/enrichment.py new file mode 100644 index 00000000..81484107 --- /dev/null +++ b/python-ecosystem/mcp-client/model/enrichment.py @@ -0,0 +1,59 @@ +from typing import Optional, Dict, List +from pydantic import BaseModel, Field + +from model.enums import RelationshipType + + +class FileContentDto(BaseModel): + """DTO representing the content of a single file retrieved from VCS.""" + path: str + content: Optional[str] = None + sizeBytes: int = 0 + skipped: bool = False + skipReason: Optional[str] = None + + +class ParsedFileMetadataDto(BaseModel): + """DTO representing parsed AST metadata for a single file.""" + path: str + language: Optional[str] = None + imports: List[str] = Field(default_factory=list) + extendsClasses: List[str] = Field(default_factory=list, alias="extends") + implementsInterfaces: List[str] = Field(default_factory=list, alias="implements") + semanticNames: List[str] = Field(default_factory=list, alias="semantic_names") + parentClass: Optional[str] = Field(default=None, alias="parent_class") + namespace: Optional[str] = None + calls: List[str] = Field(default_factory=list) + error: Optional[str] = None + + +class FileRelationshipDto(BaseModel): + """DTO representing a relationship between two files in the PR.""" + sourceFile: str + targetFile: str + relationshipType: RelationshipType + matchedOn: Optional[str] = None + strength: int = 0 + + +class EnrichmentStats(BaseModel): + """Statistics about the enrichment process.""" + totalFilesRequested: int = 0 + filesEnriched: int = 0 + filesSkipped: int = 0 + relationshipsFound: int = 0 + totalContentSizeBytes: int = 0 + processingTimeMs: int = 0 + skipReasons: Dict[str, int] = Field(default_factory=dict) + + +class PrEnrichmentDataDto(BaseModel): + """Aggregate DTO containing all file enrichment data for a PR.""" + fileContents: List[FileContentDto] = Field(default_factory=list) + fileMetadata: List[ParsedFileMetadataDto] = Field(default_factory=list) + relationships: List[FileRelationshipDto] = Field(default_factory=list) + stats: Optional[EnrichmentStats] = None + + def has_data(self) -> bool: + """Check if enrichment data is present.""" + return bool(self.fileContents) or bool(self.relationships) diff --git a/python-ecosystem/mcp-client/model/enums.py b/python-ecosystem/mcp-client/model/enums.py new file mode 100644 index 00000000..0541f03a --- /dev/null +++ b/python-ecosystem/mcp-client/model/enums.py @@ -0,0 +1,31 @@ +from enum import Enum + + +class IssueCategory(str, Enum): + """Valid issue categories for code analysis.""" + SECURITY = "SECURITY" + PERFORMANCE = "PERFORMANCE" + CODE_QUALITY = "CODE_QUALITY" + BUG_RISK = "BUG_RISK" + STYLE = "STYLE" + DOCUMENTATION = "DOCUMENTATION" + BEST_PRACTICES = "BEST_PRACTICES" + ERROR_HANDLING = "ERROR_HANDLING" + TESTING = "TESTING" + ARCHITECTURE = "ARCHITECTURE" + + +class AnalysisMode(str, Enum): + """Analysis mode for PR reviews.""" + FULL = "FULL" # Full PR diff analysis (first review or escalation) + INCREMENTAL = "INCREMENTAL" # Delta diff analysis (subsequent reviews) + + +class RelationshipType(str, Enum): + """Types of relationships between files.""" + IMPORTS = "IMPORTS" + EXTENDS = "EXTENDS" + IMPLEMENTS = "IMPLEMENTS" + CALLS = "CALLS" + SAME_PACKAGE = "SAME_PACKAGE" + REFERENCES = "REFERENCES" diff --git a/python-ecosystem/mcp-client/model/multi_stage.py b/python-ecosystem/mcp-client/model/multi_stage.py new file mode 100644 index 00000000..ff82ff77 --- /dev/null +++ b/python-ecosystem/mcp-client/model/multi_stage.py @@ -0,0 +1,87 @@ +""" +Multi-Stage Review Models. + +These models are used for the multi-stage PR review process: +- Stage 0: Planning (ReviewPlan) +- Stage 1: File-by-file review (FileReviewOutput, FileReviewBatchOutput) +- Stage 2: Cross-file analysis (CrossFileAnalysisResult) +""" + +from typing import Optional, List +from pydantic import BaseModel, Field + +from model.output_schemas import CodeReviewIssue + + +class FileReviewOutput(BaseModel): + """Stage 1 Output: Single file review result.""" + file: str + analysis_summary: str + issues: List[CodeReviewIssue] = Field(default_factory=list) + confidence: str = Field(description="Confidence level (HIGH/MEDIUM/LOW)") + note: str = Field(default="", description="Optional analysis note") + + +class FileReviewBatchOutput(BaseModel): + """Stage 1 Output: Batch of file reviews.""" + reviews: List[FileReviewOutput] = Field(description="List of review results for the files in the batch") + + +class ReviewFile(BaseModel): + """File details for review planning.""" + path: str + focus_areas: List[str] = Field(default_factory=list, description="Specific areas to focus on (SECURITY, ARCHITECTURE, etc.)") + risk_level: str = Field(default="MEDIUM", description="CRITICAL, HIGH, MEDIUM, or LOW") + estimated_issues: Optional[int] = Field(default=0) + + +class FileGroup(BaseModel): + """Group of files to be reviewed together.""" + group_id: str + priority: str = Field(description="CRITICAL, HIGH, MEDIUM, LOW") + rationale: str + files: List[ReviewFile] + + +class FileToSkip(BaseModel): + """File skipped from deep review.""" + path: str + reason: str + + +class ReviewPlan(BaseModel): + """Stage 0 Output: Plan for the review scanning.""" + analysis_summary: str + file_groups: List[FileGroup] + files_to_skip: List[FileToSkip] = Field(default_factory=list) + cross_file_concerns: List[str] = Field(default_factory=list, description="Hypotheses to verify in Stage 2") + + +class CrossFileIssue(BaseModel): + """Issue spanning multiple files (Stage 2).""" + id: str + severity: str + category: str + title: str + affected_files: List[str] + description: str + evidence: str + business_impact: str + suggestion: str + + +class DataFlowConcern(BaseModel): + """Stage 2: Data flow gap analysis.""" + flow: str + gap: str + files_involved: List[str] + severity: str + + +class CrossFileAnalysisResult(BaseModel): + """Stage 2 Output: Cross-file architectural analysis.""" + pr_risk_level: str + cross_file_issues: List[CrossFileIssue] + data_flow_concerns: List[DataFlowConcern] = Field(default_factory=list) + pr_recommendation: str + confidence: str diff --git a/python-ecosystem/mcp-client/model/output_schemas.py b/python-ecosystem/mcp-client/model/output_schemas.py new file mode 100644 index 00000000..41583db4 --- /dev/null +++ b/python-ecosystem/mcp-client/model/output_schemas.py @@ -0,0 +1,47 @@ +""" +Output Schemas for MCP Agent. + +These Pydantic models are used with MCPAgent's output_schema parameter +to ensure structured JSON output from the LLM. +""" + +from typing import Optional, List +from pydantic import BaseModel, Field + + +class CodeReviewIssue(BaseModel): + """Schema for a single code review issue.""" + # Optional issue identifier (preserve DB/client-side ids for reconciliation) + id: Optional[str] = Field(default=None, description="Optional issue id to link to existing issues") + severity: str = Field(description="Issue severity: HIGH, MEDIUM, LOW, or INFO") + category: str = Field(description="Issue category: SECURITY, PERFORMANCE, CODE_QUALITY, BUG_RISK, STYLE, DOCUMENTATION, BEST_PRACTICES, ERROR_HANDLING, TESTING, or ARCHITECTURE") + file: str = Field(description="File path where the issue is located") + line: str = Field(description="Line number or range (e.g., '42' or '42-45')") + reason: str = Field(description="Clear explanation of the issue") + suggestedFixDescription: str = Field(description="Description of the suggested fix") + suggestedFixDiff: Optional[str] = Field(default=None, description="Optional unified diff format patch for the fix") + isResolved: bool = Field(default=False, description="Whether this issue from previous analysis is resolved") + # Resolution tracking fields + resolutionExplanation: Optional[str] = Field(default=None, description="Explanation of how the issue was resolved (separate from original reason)") + resolvedInCommit: Optional[str] = Field(default=None, description="Commit hash where the issue was resolved") + # Additional fields preserved from previous issues during reconciliation + visibility: Optional[str] = Field(default=None, description="Issue visibility status") + codeSnippet: Optional[str] = Field(default=None, description="Code snippet associated with the issue") + + +class CodeReviewOutput(BaseModel): + """Schema for the complete code review output.""" + comment: str = Field(description="High-level summary of the PR analysis with key findings and recommendations") + issues: List[CodeReviewIssue] = Field(default_factory=list, description="List of identified issues in the code") + + +class SummarizeOutput(BaseModel): + """Schema for PR summarization output.""" + summary: str = Field(description="Comprehensive summary of the PR changes, purpose, and impact") + diagram: str = Field(default="", description="Visual diagram of the changes (Mermaid or ASCII format)") + diagramType: str = Field(default="MERMAID", description="Type of diagram: MERMAID or ASCII") + + +class AskOutput(BaseModel): + """Schema for ask command output.""" + answer: str = Field(description="Well-formatted markdown answer to the user's question") diff --git a/python-ecosystem/mcp-client/server/stdin_handler.py b/python-ecosystem/mcp-client/server/stdin_handler.py index c7be29fc..9817c3b8 100644 --- a/python-ecosystem/mcp-client/server/stdin_handler.py +++ b/python-ecosystem/mcp-client/server/stdin_handler.py @@ -3,8 +3,8 @@ import asyncio from typing import Optional, Dict, Any -from model.models import ReviewRequestDto -from service.review_service import ReviewService +from model.dtos import ReviewRequestDto +from service.review.review_service import ReviewService class StdinHandler: diff --git a/python-ecosystem/mcp-client/server/web_server.py b/python-ecosystem/mcp-client/server/web_server.py deleted file mode 100644 index 37a9f9f7..00000000 --- a/python-ecosystem/mcp-client/server/web_server.py +++ /dev/null @@ -1,309 +0,0 @@ -import os -import json -import shutil -import tempfile -import tarfile -import asyncio -from typing import Dict, Any -from fastapi import FastAPI, Request, UploadFile, File, Form -from starlette.responses import StreamingResponse - -from model.models import ( - ReviewRequestDto, ReviewResponseDto, - SummarizeRequestDto, SummarizeResponseDto, - AskRequestDto, AskResponseDto -) -from service.review_service import ReviewService -from service.command_service import CommandService -from utils.response_parser import ResponseParser - - -def create_app(): - """Create and configure FastAPI application.""" - app = FastAPI(title="codecrow-mcp-client") - review_service = ReviewService() - command_service = CommandService() - - @app.post("/review", response_model=ReviewResponseDto) - async def review_endpoint(req: ReviewRequestDto, request: Request): - """ - HTTP endpoint to accept review requests from the pipeline agent. - - Behavior: - - If the client requests streaming via header `Accept: application/x-ndjson`, - the endpoint will return a StreamingResponse that yields NDJSON events as they occur. - - Otherwise it preserves the original behavior and returns a single ReviewResponseDto JSON body. - """ - try: - wants_stream = _wants_streaming(request) - - if not wants_stream: - # Non-streaming (legacy) behavior - result = await review_service.process_review_request(req) - return ReviewResponseDto( - result=result.get("result"), - error=result.get("error") - ) - - # Streaming behavior - async def event_stream(): - queue = asyncio.Queue() - - # Emit initial queued status - yield _json_event({"type": "status", "state": "queued", "message": "request received"}) - - # Event callback to capture service events - def event_callback(event: Dict[str, Any]): - try: - queue.put_nowait(event) - except asyncio.QueueFull: - pass # Skip if queue is full - - # Run processing in background - async def runner(): - try: - result = await review_service.process_review_request( - req, - event_callback=event_callback - ) - # Emit final event with result - final_event = { - "type": "final", - "result": result.get("result") - } - await queue.put(final_event) - except Exception as e: - await queue.put({ - "type": "error", - "message": str(e) - }) - - task = asyncio.create_task(runner()) - - # Drain queue and yield events - async for event in _drain_queue_until_final(queue, task): - yield _json_event(event) - - return StreamingResponse(event_stream(), media_type="application/x-ndjson") - - except Exception as e: - error_response = ResponseParser.create_error_response( - "HTTP request processing failed", str(e) - ) - return ReviewResponseDto(result=error_response) - - @app.post("/review/summarize", response_model=SummarizeResponseDto) - async def summarize_endpoint(req: SummarizeRequestDto, request: Request): - """ - HTTP endpoint to process /codecrow summarize command. - - Generates a comprehensive PR summary with: - - Overview of changes - - Key files modified - - Impact analysis - - Architecture diagram (Mermaid or ASCII) - """ - try: - wants_stream = _wants_streaming(request) - - if not wants_stream: - # Non-streaming behavior - result = await command_service.process_summarize(req) - return SummarizeResponseDto( - summary=result.get("summary"), - diagram=result.get("diagram"), - diagramType=result.get("diagramType", "MERMAID"), - error=result.get("error") - ) - - # Streaming behavior - async def event_stream(): - queue = asyncio.Queue() - - yield _json_event({"type": "status", "state": "queued", "message": "summarize request received"}) - - def event_callback(event: Dict[str, Any]): - try: - queue.put_nowait(event) - except asyncio.QueueFull: - pass - - async def runner(): - try: - result = await command_service.process_summarize(req, event_callback=event_callback) - await queue.put({ - "type": "final", - "result": result - }) - except Exception as e: - await queue.put({"type": "error", "message": str(e)}) - - task = asyncio.create_task(runner()) - - async for event in _drain_queue_until_final(queue, task): - yield _json_event(event) - - return StreamingResponse(event_stream(), media_type="application/x-ndjson") - - except Exception as e: - return SummarizeResponseDto(error=f"Summarize failed: {str(e)}") - - @app.post("/review/ask", response_model=AskResponseDto) - async def ask_endpoint(req: AskRequestDto, request: Request): - """ - HTTP endpoint to process /codecrow ask command. - - Answers questions about: - - Specific issues - - PR changes - - Codebase (using RAG) - - Analysis results - """ - try: - wants_stream = _wants_streaming(request) - - if not wants_stream: - # Non-streaming behavior - result = await command_service.process_ask(req) - return AskResponseDto( - answer=result.get("answer"), - error=result.get("error") - ) - - # Streaming behavior - async def event_stream(): - queue = asyncio.Queue() - - yield _json_event({"type": "status", "state": "queued", "message": "ask request received"}) - - def event_callback(event: Dict[str, Any]): - try: - queue.put_nowait(event) - except asyncio.QueueFull: - pass - - async def runner(): - try: - result = await command_service.process_ask(req, event_callback=event_callback) - await queue.put({ - "type": "final", - "result": result - }) - except Exception as e: - await queue.put({"type": "error", "message": str(e)}) - - task = asyncio.create_task(runner()) - - async for event in _drain_queue_until_final(queue, task): - yield _json_event(event) - - return StreamingResponse(event_stream(), media_type="application/x-ndjson") - - except Exception as e: - return AskResponseDto(error=f"Ask failed: {str(e)}") - - @app.get("/health") - def health(): - """Health check endpoint.""" - return {"status": "ok"} - - return app - - -def _wants_streaming(request: Request) -> bool: - """Check if client wants streaming response.""" - accept_header = request.headers.get("accept", "") - return "application/x-ndjson" in accept_header.lower() - - -def _json_event(event: Dict[str, Any]) -> str: - """Serialize event to NDJSON line.""" - return json.dumps(event) + "\n" - - -async def _extract_archive(archive_path: str, extract_dir: str) -> bool: - """ - Extract tar.gz or zip archive. - Returns True if successful, False otherwise. - """ - os.makedirs(extract_dir, exist_ok=True) - - # Try tar first - try: - with tarfile.open(archive_path, "r:*") as tar: - tar.extractall(path=extract_dir) - return True - except tarfile.TarError: - pass - - # Try zip - try: - shutil.unpack_archive(archive_path, extract_dir) - return True - except Exception: - pass - - return False - - -async def _drain_queue_until_final(queue: asyncio.Queue, task: asyncio.Task): - """ - Drain events from queue until we see a final/error event and task is done. - Yields each event as it arrives. - """ - seen_final = False - - while True: - try: - # Wait for event with timeout to prevent hanging - event = await asyncio.wait_for(queue.get(), timeout=1.0) - yield event - - # Check if this is a terminal event - event_type = event.get("type") - if event_type in ("final", "error"): - seen_final = True - # Continue draining in case there are more events - - # If we've seen final and task is done, check for remaining events then exit - if seen_final and task.done(): - # Give a moment for any last events - await asyncio.sleep(0.1) - try: - while True: - event = queue.get_nowait() - yield event - except asyncio.QueueEmpty: - break - break - - except asyncio.TimeoutError: - # No event available, check if task is done - if task.done(): - # Task finished, drain any remaining events - try: - while True: - event = queue.get_nowait() - yield event - if event.get("type") in ("final", "error"): - seen_final = True - except asyncio.QueueEmpty: - pass - - # If we saw a final event or no more events, we're done - if seen_final or task.done(): - break - # Otherwise continue waiting for events - - -def run_http_server(host: str = "0.0.0.0", port: int = 8000): - """Run the FastAPI application.""" - app = create_app() - import uvicorn - uvicorn.run(app, host=host, port=port, log_level="info", timeout_keep_alive=300) - - -if __name__ == "__main__": - host = os.environ.get("AI_CLIENT_HOST", "0.0.0.0") - port = int(os.environ.get("AI_CLIENT_PORT", "8000")) - run_http_server(host=host, port=port) \ No newline at end of file diff --git a/python-ecosystem/mcp-client/service/__init__.py b/python-ecosystem/mcp-client/service/__init__.py new file mode 100644 index 00000000..6a893b0b --- /dev/null +++ b/python-ecosystem/mcp-client/service/__init__.py @@ -0,0 +1,44 @@ +""" +Service Package. + +Contains business logic services organized into subpackages: +- review: Code review functionality (ReviewService, orchestrator, issue processing) +- rag: RAG pipeline client and reranking +- command: Command handling (summarize, ask) +""" + +# Re-export from subpackages for backward compatibility +from service.review import ( + ReviewService, + MultiStageReviewOrchestrator, + IssuePostProcessor, + IssueDeduplicator, + post_process_analysis_result, + restore_missing_diffs_from_previous, +) +from service.rag import ( + RagClient, + RAG_MIN_RELEVANCE_SCORE, + RAG_DEFAULT_TOP_K, + LLMReranker, + RerankResult, +) +from service.command import CommandService + +__all__ = [ + # Review + "ReviewService", + "MultiStageReviewOrchestrator", + "IssuePostProcessor", + "IssueDeduplicator", + "post_process_analysis_result", + "restore_missing_diffs_from_previous", + # RAG + "RagClient", + "RAG_MIN_RELEVANCE_SCORE", + "RAG_DEFAULT_TOP_K", + "LLMReranker", + "RerankResult", + # Command + "CommandService", +] diff --git a/python-ecosystem/mcp-client/service/command/__init__.py b/python-ecosystem/mcp-client/service/command/__init__.py new file mode 100644 index 00000000..3019afae --- /dev/null +++ b/python-ecosystem/mcp-client/service/command/__init__.py @@ -0,0 +1,8 @@ +""" +Command Service Package. + +Contains the CommandService for handling CodeCrow commands (summarize, ask). +""" +from service.command.command_service import CommandService + +__all__ = ["CommandService"] diff --git a/python-ecosystem/mcp-client/service/command_service.py b/python-ecosystem/mcp-client/service/command/command_service.py similarity index 93% rename from python-ecosystem/mcp-client/service/command_service.py rename to python-ecosystem/mcp-client/service/command/command_service.py index 9de4898b..f8743e6a 100644 --- a/python-ecosystem/mcp-client/service/command_service.py +++ b/python-ecosystem/mcp-client/service/command/command_service.py @@ -10,10 +10,11 @@ from mcp_use import MCPAgent, MCPClient from langchain_core.agents import AgentAction -from model.models import SummarizeRequestDto, AskRequestDto, SummarizeOutput, AskOutput +from model.dtos import SummarizeRequestDto, AskRequestDto +from model.output_schemas import SummarizeOutput, AskOutput from utils.mcp_config import MCPConfigBuilder from llm.llm_factory import LLMFactory -from service.rag_client import RagClient +from service.rag.rag_client import RagClient from utils.error_sanitizer import sanitize_error_for_display, create_user_friendly_error logger = logging.getLogger(__name__) @@ -27,7 +28,7 @@ class CommandService: MAX_STEPS_ASK = 40 def __init__(self): - load_dotenv() + load_dotenv(interpolate=False) self.default_jar_path = os.environ.get( "MCP_SERVER_JAR", "/app/codecrow-vcs-mcp-1.0.jar" @@ -56,11 +57,12 @@ async def process_summarize( return {"error": error_msg} try: - self._emit_event(event_callback, { - "type": "status", - "state": "started", - "message": "Starting PR summarization" - }) + async with asyncio.timeout(300): # 5-minute ceiling for commands + self._emit_event(event_callback, { + "type": "status", + "state": "started", + "message": "Starting PR summarization" + }) # Build configuration jvm_props = self._build_jvm_props_for_summarize(request) @@ -92,13 +94,20 @@ async def process_summarize( # TODO: Mermaid diagrams disabled for now - AI-generated Mermaid often has syntax errors # that fail to render on GitHub. Using ASCII diagrams until we add validation/fixing. # Original: supports_mermaid=request.supportsMermaid - result = await self._execute_summarize( - llm=llm, - client=client, - prompt=prompt, - supports_mermaid=False, # Mermaid disabled - always use ASCII - event_callback=event_callback - ) + try: + result = await self._execute_summarize( + llm=llm, + client=client, + prompt=prompt, + supports_mermaid=False, # Mermaid disabled - always use ASCII + event_callback=event_callback + ) + finally: + # Always close MCP sessions to release JVM subprocesses + try: + await client.close_all_sessions() + except Exception as close_err: + logger.warning(f"Error closing MCP sessions: {close_err}") self._emit_event(event_callback, { "type": "final", @@ -107,6 +116,12 @@ async def process_summarize( return result + except TimeoutError: + timeout_msg = "Summarize command timed out after 300 seconds" + logger.error(timeout_msg) + self._emit_event(event_callback, {"type": "error", "message": timeout_msg}) + return {"error": timeout_msg} + except Exception as e: logger.error(f"Summarize failed: {str(e)}", exc_info=True) sanitized_msg = create_user_friendly_error(e) @@ -141,11 +156,12 @@ async def process_ask( ) try: - self._emit_event(event_callback, { - "type": "status", - "state": "started", - "message": "Processing your question" - }) + async with asyncio.timeout(300): # 5-minute ceiling for commands + self._emit_event(event_callback, { + "type": "status", + "state": "started", + "message": "Processing your question" + }) # Build configuration with both VCS and Platform MCP servers jvm_props = self._build_jvm_props_for_ask(request) @@ -188,12 +204,19 @@ async def process_ask( }) # Execute with MCP agent - result = await self._execute_ask( - llm=llm, - client=client, - prompt=prompt, - event_callback=event_callback - ) + try: + result = await self._execute_ask( + llm=llm, + client=client, + prompt=prompt, + event_callback=event_callback + ) + finally: + # Always close MCP sessions to release JVM subprocesses + try: + await client.close_all_sessions() + except Exception as close_err: + logger.warning(f"Error closing MCP sessions: {close_err}") self._emit_event(event_callback, { "type": "final", @@ -202,6 +225,12 @@ async def process_ask( return result + except TimeoutError: + timeout_msg = "Ask command timed out after 300 seconds" + logger.error(timeout_msg) + self._emit_event(event_callback, {"type": "error", "message": timeout_msg}) + return {"error": timeout_msg} + except Exception as e: logger.error(f"Ask failed: {str(e)}", exc_info=True) sanitized_msg = create_user_friendly_error(e) diff --git a/python-ecosystem/mcp-client/service/multi_stage_orchestrator.py b/python-ecosystem/mcp-client/service/multi_stage_orchestrator.py deleted file mode 100644 index 70fa3cd7..00000000 --- a/python-ecosystem/mcp-client/service/multi_stage_orchestrator.py +++ /dev/null @@ -1,1117 +0,0 @@ -import logging -import asyncio -import json -from typing import Dict, Any, List, Optional, Callable - -from model.models import ( - ReviewRequestDto, - ReviewPlan, - CodeReviewOutput, - CodeReviewIssue, - ReviewFile, - FileGroup, - CrossFileAnalysisResult, - FileReviewBatchOutput -) -from utils.prompts.prompt_builder import PromptBuilder -from utils.diff_processor import ProcessedDiff, DiffProcessor -from mcp_use import MCPAgent - -logger = logging.getLogger(__name__) - - -def extract_llm_response_text(response: Any) -> str: - """ - Extract text content from LLM response, handling different response formats. - Some LLM providers return content as a list of objects instead of a string. - """ - if hasattr(response, 'content'): - content = response.content - if isinstance(content, list): - # Handle list content (e.g., from Gemini or other providers) - text_parts = [] - for item in content: - if isinstance(item, str): - text_parts.append(item) - elif isinstance(item, dict): - if 'text' in item: - text_parts.append(item['text']) - elif 'content' in item: - text_parts.append(item['content']) - elif hasattr(item, 'text'): - text_parts.append(item.text) - return "".join(text_parts) - return str(content) - return str(response) - - -# Prevent duplicate logs from mcp_use -mcp_logger = logging.getLogger("mcp_use") -mcp_logger.propagate = False -if not mcp_logger.handlers: - handler = logging.StreamHandler() - formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') - handler.setFormatter(formatter) - mcp_logger.addHandler(handler) - -class RecursiveMCPAgent(MCPAgent): - """ - Subclass of MCPAgent that enforces a higher recursion limit on the internal agent executor. - """ - def __init__(self, *args, recursion_limit: int = 50, **kwargs): - self._custom_recursion_limit = recursion_limit - super().__init__(*args, **kwargs) - - async def stream(self, *args, **kwargs): - """ - Override stream to ensure recursion_limit is applied. - """ - # Ensure the executor exists - if self._agent_executor is None: - await self.initialize() - - # Patch the executor's astream if not already patched - executor = self._agent_executor - if executor and not getattr(executor, "_is_patched_recursion", False): - original_astream = executor.astream - limit = self._custom_recursion_limit - - async def patched_astream(input_data, config=None, **astream_kwargs): - if config is None: - config = {} - config["recursion_limit"] = limit - async for chunk in original_astream(input_data, config=config, **astream_kwargs): - yield chunk - - executor.astream = patched_astream - executor._is_patched_recursion = True - logger.info(f"RecursiveMCPAgent: Patched recursion limit to {limit}") - - # Call parent stream - async for item in super().stream(*args, **kwargs): - yield item - -class MultiStageReviewOrchestrator: - """ - Orchestrates the 4-stage AI code review pipeline: - Stage 0: Planning & Prioritization - Stage 1: Parallel File Review - Stage 2: Cross-File & Architectural Analysis - Stage 3: Aggregation & Final Report - """ - - def __init__( - self, - llm, - mcp_client, - rag_client=None, - event_callback: Optional[Callable[[Dict], None]] = None - ): - self.llm = llm - self.client = mcp_client - self.rag_client = rag_client - self.event_callback = event_callback - self.max_parallel_stage_1 = 5 # Limit parallel execution to avoid rate limits - - async def execute_branch_analysis(self, prompt: str) -> Dict[str, Any]: - """ - Execute a single-pass branch analysis using the provided prompt. - """ - self._emit_status("branch_analysis_started", "Starting Branch Analysis & Reconciliation...") - - agent = RecursiveMCPAgent( - llm=self.llm, - client=self.client, - additional_instructions=PromptBuilder.get_additional_instructions() - ) - - - try: - final_text = "" - # Branch analysis expects standard CodeReviewOutput - async for item in agent.stream(prompt, max_steps=15, output_schema=CodeReviewOutput): - if isinstance(item, CodeReviewOutput): - # Convert to dict format expected by service - issues = [i.model_dump() for i in item.issues] if item.issues else [] - return { - "issues": issues, - "comment": item.comment or "Branch analysis completed." - } - - if isinstance(item, str): - final_text = item - - # If stream finished without object, try parsing text - if final_text: - data = await self._parse_response(final_text, CodeReviewOutput) - issues = [i.model_dump() for i in data.issues] if data.issues else [] - return { - "issues": issues, - "comment": data.comment or "Branch analysis completed." - } - - return {"issues": [], "comment": "No issues found."} - - except Exception as e: - logger.error(f"Branch analysis failed: {e}", exc_info=True) - self._emit_error(str(e)) - raise - - async def orchestrate_review( - self, - request: ReviewRequestDto, - rag_context: Optional[Dict[str, Any]] = None, - processed_diff: Optional[ProcessedDiff] = None - ) -> Dict[str, Any]: - """ - Main entry point for the multi-stage review. - Supports both FULL (initial review) and INCREMENTAL (follow-up review) modes. - The same pipeline is used, but with incremental-aware prompts and issue reconciliation. - """ - # Determine if this is an incremental review - is_incremental = ( - request.analysisMode == "INCREMENTAL" - and request.deltaDiff - ) - - if is_incremental: - logger.info(f"INCREMENTAL mode: reviewing delta diff, {len(request.previousCodeAnalysisIssues or [])} previous issues to reconcile") - else: - logger.info("FULL mode: initial PR review") - - try: - # === STAGE 0: Planning === - self._emit_status("stage_0_started", "Stage 0: Planning & Prioritization...") - review_plan = await self._execute_stage_0_planning(request, is_incremental) - - # Validate and fix the plan to ensure all files are included - review_plan = self._ensure_all_files_planned(review_plan, request.changedFiles or []) - self._emit_progress(10, "Stage 0 Complete: Review plan created") - - # === STAGE 1: File Reviews === - self._emit_status("stage_1_started", f"Stage 1: Analyzing {self._count_files(review_plan)} files...") - file_issues = await self._execute_stage_1_file_reviews( - request, review_plan, rag_context, processed_diff, is_incremental - ) - self._emit_progress(60, f"Stage 1 Complete: {len(file_issues)} issues found across files") - - # === STAGE 1.5: Issue Reconciliation (INCREMENTAL only) === - if is_incremental and request.previousCodeAnalysisIssues: - self._emit_status("reconciliation_started", "Reconciling previous issues...") - file_issues = await self._reconcile_previous_issues( - request, file_issues, processed_diff - ) - self._emit_progress(70, f"Reconciliation Complete: {len(file_issues)} total issues after reconciliation") - - # === STAGE 2: Cross-File Analysis === - self._emit_status("stage_2_started", "Stage 2: Analyzing cross-file patterns...") - cross_file_results = await self._execute_stage_2_cross_file(request, file_issues, review_plan) - self._emit_progress(85, "Stage 2 Complete: Cross-file analysis finished") - - # === STAGE 3: Aggregation === - self._emit_status("stage_3_started", "Stage 3: Generating final report...") - final_report = await self._execute_stage_3_aggregation( - request, review_plan, file_issues, cross_file_results, is_incremental - ) - self._emit_progress(100, "Stage 3 Complete: Report generated") - - # Return structure compatible with existing response expected by frontend/controller - return { - "comment": final_report, - "issues": [issue.model_dump() for issue in file_issues], - } - - except Exception as e: - logger.error(f"Multi-stage review failed: {e}", exc_info=True) - self._emit_error(str(e)) - raise - - async def _reconcile_previous_issues( - self, - request: ReviewRequestDto, - new_issues: List[CodeReviewIssue], - processed_diff: Optional[ProcessedDiff] = None - ) -> List[CodeReviewIssue]: - """ - Reconcile previous issues with new findings in incremental mode. - - Mark resolved issues as isResolved=true - - Update line numbers for persisting issues - - Merge with new issues found in delta diff - """ - if not request.previousCodeAnalysisIssues: - return new_issues - - logger.info(f"Reconciling {len(request.previousCodeAnalysisIssues)} previous issues with {len(new_issues)} new issues") - - # Get the delta diff content to check what files/lines changed - delta_diff = request.deltaDiff or "" - - # Build a set of files that changed in the delta - changed_files_in_delta = set() - if processed_diff: - for f in processed_diff.files: - changed_files_in_delta.add(f.path) - - reconciled_issues = list(new_issues) # Start with new issues - - # Process each previous issue - for prev_issue in request.previousCodeAnalysisIssues: - # Convert to dict if needed - if hasattr(prev_issue, 'model_dump'): - prev_data = prev_issue.model_dump() - else: - prev_data = prev_issue if isinstance(prev_issue, dict) else vars(prev_issue) - - # Debug log to verify field mapping - logger.debug(f"Previous issue data: reason={prev_data.get('reason')}, " - f"suggestedFixDescription={prev_data.get('suggestedFixDescription')}, " - f"suggestedFixDiff={prev_data.get('suggestedFixDiff')[:50] if prev_data.get('suggestedFixDiff') else None}") - - file_path = prev_data.get('file', prev_data.get('filePath', '')) - issue_id = prev_data.get('id') - - # Check if this issue was already found in new issues (by file+line or ID) - already_reported = False - for new_issue in new_issues: - new_data = new_issue.model_dump() if hasattr(new_issue, 'model_dump') else new_issue - if (new_data.get('file') == file_path and - str(new_data.get('line')) == str(prev_data.get('line', prev_data.get('lineNumber')))): - already_reported = True - break - if issue_id and new_data.get('id') == issue_id: - already_reported = True - break - - if already_reported: - continue # Already in new issues, skip - - # Check if the file was modified in delta diff - file_in_delta = any(file_path.endswith(f) or f.endswith(file_path) for f in changed_files_in_delta) - - # If file wasn't touched in delta, issue persists unchanged - # If file was touched, we need to check if the specific line was modified - is_resolved = False - if file_in_delta: - # Simple heuristic: if file changed and we didn't re-report this issue, - # it might be resolved. But we should be conservative here. - # For now, we'll re-report it as persisting unless LLM marked it resolved - pass - - # Preserve all original issue data - just pass through as CodeReviewIssue - # Field mapping from Java DTO: - # reason (or title for legacy) -> reason - # severity (uppercase) -> severity - # category (or issueCategory) -> category - # file -> file - # line -> line - # suggestedFixDescription -> suggestedFixDescription - # suggestedFixDiff -> suggestedFixDiff - persisting_issue = CodeReviewIssue( - id=str(issue_id) if issue_id else None, - severity=(prev_data.get('severity') or prev_data.get('issueSeverity') or 'MEDIUM').upper(), - category=prev_data.get('category') or prev_data.get('issueCategory') or prev_data.get('type') or 'CODE_QUALITY', - file=file_path or prev_data.get('file') or prev_data.get('filePath') or 'unknown', - line=str(prev_data.get('line') or prev_data.get('lineNumber') or '1'), - reason=prev_data.get('reason') or prev_data.get('title') or prev_data.get('description') or '', - suggestedFixDescription=prev_data.get('suggestedFixDescription') or prev_data.get('suggestedFix') or '', - suggestedFixDiff=prev_data.get('suggestedFixDiff') or None, - isResolved=is_resolved, - visibility=prev_data.get('visibility'), - codeSnippet=prev_data.get('codeSnippet') - ) - reconciled_issues.append(persisting_issue) - - logger.info(f"Reconciliation complete: {len(reconciled_issues)} total issues") - return reconciled_issues - - def _issue_matches_files(self, issue: Any, file_paths: List[str]) -> bool: - """Check if an issue is related to any of the given file paths.""" - if hasattr(issue, 'model_dump'): - issue_data = issue.model_dump() - elif isinstance(issue, dict): - issue_data = issue - else: - issue_data = vars(issue) if hasattr(issue, '__dict__') else {} - - issue_file = issue_data.get('file', issue_data.get('filePath', '')) - - for fp in file_paths: - if issue_file == fp or issue_file.endswith('/' + fp) or fp.endswith('/' + issue_file): - return True - # Also check basename match - if issue_file.split('/')[-1] == fp.split('/')[-1]: - return True - return False - - def _format_previous_issues_for_batch(self, issues: List[Any]) -> str: - """Format previous issues for inclusion in batch prompt.""" - if not issues: - return "" - - lines = ["=== PREVIOUS ISSUES IN THESE FILES (check if resolved) ==="] - for issue in issues: - if hasattr(issue, 'model_dump'): - data = issue.model_dump() - elif isinstance(issue, dict): - data = issue - else: - data = vars(issue) if hasattr(issue, '__dict__') else {} - - issue_id = data.get('id', 'unknown') - severity = data.get('severity', 'MEDIUM') - file_path = data.get('file', data.get('filePath', 'unknown')) - line = data.get('line', data.get('lineNumber', '?')) - reason = data.get('reason', data.get('description', 'No description')) - - lines.append(f"[ID:{issue_id}] {severity} @ {file_path}:{line}") - lines.append(f" Issue: {reason}") - lines.append("") - - lines.append("Mark these as 'isResolved: true' if fixed in the diff above.") - lines.append("=== END PREVIOUS ISSUES ===") - return "\n".join(lines) - - def _get_diff_snippets_for_batch( - self, - all_diff_snippets: List[str], - batch_file_paths: List[str] - ) -> List[str]: - """ - Filter diff snippets to only include those relevant to the batch files. - - The diffSnippets from ReviewRequestDto are extracted by Java DiffParser.extractDiffSnippets() - which creates snippets in format that may include file paths in diff headers. - We filter to get only snippets relevant to the current batch for targeted RAG queries. - """ - if not all_diff_snippets or not batch_file_paths: - return [] - - batch_snippets = [] - batch_file_names = {path.split('/')[-1] for path in batch_file_paths} - - for snippet in all_diff_snippets: - # Check if snippet relates to any file in the batch - # Diff snippets typically contain file paths in headers like "diff --git a/path b/path" - # or "--- a/path" or "+++ b/path" - snippet_lower = snippet.lower() - - for file_path in batch_file_paths: - if file_path.lower() in snippet_lower: - batch_snippets.append(snippet) - break - else: - # Also check by filename only (for cases where paths differ) - for file_name in batch_file_names: - if file_name.lower() in snippet_lower: - batch_snippets.append(snippet) - break - - logger.debug(f"Filtered {len(batch_snippets)} diff snippets for batch from {len(all_diff_snippets)} total") - return batch_snippets - - async def _execute_stage_0_planning(self, request: ReviewRequestDto, is_incremental: bool = False) -> ReviewPlan: - """ - Stage 0: Analyze metadata and generate a review plan. - Uses structured output for reliable JSON parsing. - """ - # Prepare context for prompt - changed_files_summary = [] - if request.changedFiles: - for f in request.changedFiles: - changed_files_summary.append({ - "path": f, - "type": "MODIFIED", - "lines_added": "?", - "lines_deleted": "?" - }) - - prompt = PromptBuilder.build_stage_0_planning_prompt( - repo_slug=request.projectVcsRepoSlug, - pr_id=str(request.pullRequestId), - pr_title=request.prTitle or "", - author="Unknown", - branch_name="source-branch", - target_branch=request.targetBranchName or "main", - commit_hash=request.commitHash or "HEAD", - changed_files_json=json.dumps(changed_files_summary, indent=2) - ) - - # Stage 0 uses direct LLM call (no tools needed - all metadata is provided) - try: - structured_llm = self.llm.with_structured_output(ReviewPlan) - result = await structured_llm.ainvoke(prompt) - if result: - logger.info("Stage 0 planning completed with structured output") - return result - except Exception as e: - logger.warning(f"Structured output failed for Stage 0: {e}") - - # Fallback to manual parsing - try: - response = await self.llm.ainvoke(prompt) - content = extract_llm_response_text(response) - return await self._parse_response(content, ReviewPlan) - except Exception as e: - logger.error(f"Stage 0 planning failed: {e}") - raise ValueError(f"Stage 0 planning failed: {e}") - - def _chunk_files(self, file_groups: List[Any], max_files_per_batch: int = 5) -> List[List[Dict[str, Any]]]: - """Flatten file groups and chunk into batches.""" - all_files = [] - for group in file_groups: - for f in group.files: - # Attach priority context for the review - all_files.append({ - "file": f, - "priority": group.priority - }) - - return [all_files[i:i + max_files_per_batch] for i in range(0, len(all_files), max_files_per_batch)] - - async def _execute_stage_1_file_reviews( - self, - request: ReviewRequestDto, - plan: ReviewPlan, - rag_context: Optional[Dict[str, Any]] = None, - processed_diff: Optional[ProcessedDiff] = None, - is_incremental: bool = False - ) -> List[CodeReviewIssue]: - """ - Stage 1: Execute batch file reviews with per-batch RAG context. - """ - # Use smaller batches (3 files max) for better RAG relevance and review quality - batches = self._chunk_files(plan.file_groups, max_files_per_batch=3) - - total_files = sum(len(batch) for batch in batches) - logger.info(f"Stage 1: Processing {total_files} files in {len(batches)} batches (max 3 files/batch)") - - # Process batches with controlled parallelism - all_issues = [] - batch_results = [] - - # Process in waves to avoid rate limits - for wave_start in range(0, len(batches), self.max_parallel_stage_1): - wave_end = min(wave_start + self.max_parallel_stage_1, len(batches)) - wave_batches = batches[wave_start:wave_end] - - logger.info(f"Stage 1: Processing wave {wave_start // self.max_parallel_stage_1 + 1}, " - f"batches {wave_start + 1}-{wave_end} of {len(batches)}") - - tasks = [] - for batch_idx, batch in enumerate(wave_batches, start=wave_start + 1): - batch_paths = [item["file"].path for item in batch] - logger.debug(f"Batch {batch_idx}: {batch_paths}") - tasks.append(self._review_file_batch(request, batch, processed_diff, is_incremental)) - - results = await asyncio.gather(*tasks, return_exceptions=True) - - for idx, res in enumerate(results): - batch_num = wave_start + idx + 1 - if isinstance(res, Exception): - logger.error(f"Error reviewing batch {batch_num}: {res}") - elif res: - logger.info(f"Batch {batch_num} completed: {len(res)} issues found") - all_issues.extend(res) - else: - logger.info(f"Batch {batch_num} completed: no issues found") - - # Update progress - progress = 10 + int((wave_end / len(batches)) * 50) - self._emit_progress(progress, f"Stage 1: Reviewed {wave_end}/{len(batches)} batches") - - logger.info(f"Stage 1 Complete: {len(all_issues)} issues found across {total_files} files") - return all_issues - - async def _review_file_batch( - self, - request: ReviewRequestDto, - batch_items: List[Dict[str, Any]], - processed_diff: Optional[ProcessedDiff] = None, - is_incremental: bool = False - ) -> List[CodeReviewIssue]: - """ - Review a batch of files in a single LLM call with per-batch RAG context. - In incremental mode, uses delta diff and focuses on new changes only. - """ - batch_files_data = [] - batch_file_paths = [] - project_rules = "1. No hardcoded secrets.\n2. Use dependency injection.\n3. Verify all inputs." - - # For incremental mode, use deltaDiff instead of full diff - diff_source = None - if is_incremental and request.deltaDiff: - # Parse delta diff to extract per-file diffs - diff_source = DiffProcessor().process(request.deltaDiff) if request.deltaDiff else None - else: - diff_source = processed_diff - - # Collect file paths and diffs for this batch - for item in batch_items: - file_info = item["file"] - batch_file_paths.append(file_info.path) - - # Extract diff from the appropriate source (delta for incremental, full for initial) - file_diff = "" - if diff_source: - for f in diff_source.files: - if f.path == file_info.path or f.path.endswith("/" + file_info.path): - file_diff = f.content - break - - batch_files_data.append({ - "path": file_info.path, - "type": "MODIFIED", - "focus_areas": file_info.focus_areas, - "old_code": "", - "diff": file_diff or "(Diff unavailable)", - "is_incremental": is_incremental # Pass mode to prompt builder - }) - - # Filter pre-computed diff snippets for files in this batch (for RAG query) - # The diffSnippets from ReviewRequestDto are already properly extracted by Java DiffParser - batch_diff_snippets = self._get_diff_snippets_for_batch( - request.diffSnippets or [], - batch_file_paths - ) - - # Fetch per-batch RAG context (targeted to these specific files) - rag_context_text = "" - if self.rag_client and self.rag_client.enabled: - try: - rag_response = await self.rag_client.get_pr_context( - workspace=request.projectWorkspace, - project=request.projectVcsRepoSlug, - branch=request.targetBranchName, - changed_files=batch_file_paths, - diff_snippets=batch_diff_snippets, - pr_title=request.prTitle, - top_k=5, # Focused context for this batch - min_relevance_score=0.75 # Higher threshold for per-batch - ) - # Pass ALL changed files from the PR to filter out stale context - # RAG returns context from main branch, so files being modified in PR are stale - rag_context_text = self._format_rag_context( - rag_response.get("context", {}), - set(batch_file_paths), - pr_changed_files=request.changedFiles - ) - logger.debug(f"Batch RAG context retrieved for {batch_file_paths}") - except Exception as e: - logger.warning(f"Failed to fetch per-batch RAG context: {e}") - - # For incremental mode, filter previous issues relevant to this batch - previous_issues_for_batch = "" - if is_incremental and request.previousCodeAnalysisIssues: - relevant_prev_issues = [ - issue for issue in request.previousCodeAnalysisIssues - if self._issue_matches_files(issue, batch_file_paths) - ] - if relevant_prev_issues: - previous_issues_for_batch = self._format_previous_issues_for_batch(relevant_prev_issues) - - # Build ONE prompt for the batch - prompt = PromptBuilder.build_stage_1_batch_prompt( - files=batch_files_data, - priority=batch_items[0]["priority"] if batch_items else "MEDIUM", - project_rules=project_rules, - rag_context=rag_context_text, - is_incremental=is_incremental, - previous_issues=previous_issues_for_batch - ) - - # Stage 1 uses direct LLM call (no tools needed - diff is already provided) - try: - # Try structured output first - structured_llm = self.llm.with_structured_output(FileReviewBatchOutput) - result = await structured_llm.ainvoke(prompt) - if result: - all_batch_issues = [] - for review in result.reviews: - all_batch_issues.extend(review.issues) - return all_batch_issues - except Exception as e: - logger.warning(f"Structured output failed for Stage 1 batch: {e}") - - # Fallback to manual parsing - try: - response = await self.llm.ainvoke(prompt) - content = extract_llm_response_text(response) - data = await self._parse_response(content, FileReviewBatchOutput) - all_batch_issues = [] - for review in data.reviews: - all_batch_issues.extend(review.issues) - return all_batch_issues - except Exception as parse_err: - logger.error(f"Batch review failed: {parse_err}") - return [] - - return [] - - async def _execute_stage_2_cross_file( - self, - request: ReviewRequestDto, - stage_1_issues: List[CodeReviewIssue], - plan: ReviewPlan - ) -> CrossFileAnalysisResult: - """ - Stage 2: Cross-file analysis. - """ - # Serialize Stage 1 findings - issues_json = json.dumps([i.model_dump() for i in stage_1_issues], indent=2) - - prompt = PromptBuilder.build_stage_2_cross_file_prompt( - repo_slug=request.projectVcsRepoSlug, - pr_title=request.prTitle or "", - commit_hash=request.commitHash or "HEAD", - stage_1_findings_json=issues_json, - architecture_context="(Architecture context from MCP or knowledge base)", - migrations="(Migration scripts found in PR)", - cross_file_concerns=plan.cross_file_concerns - ) - - # Stage 2 uses direct LLM call (no tools needed - all data is provided from Stage 1) - try: - structured_llm = self.llm.with_structured_output(CrossFileAnalysisResult) - result = await structured_llm.ainvoke(prompt) - if result: - logger.info("Stage 2 cross-file analysis completed with structured output") - return result - except Exception as e: - logger.warning(f"Structured output failed for Stage 2: {e}") - - # Fallback to manual parsing - try: - response = await self.llm.ainvoke(prompt) - content = extract_llm_response_text(response) - return await self._parse_response(content, CrossFileAnalysisResult) - except Exception as e: - logger.error(f"Stage 2 cross-file analysis failed: {e}") - raise - - async def _execute_stage_3_aggregation( - self, - request: ReviewRequestDto, - plan: ReviewPlan, - stage_1_issues: List[CodeReviewIssue], - stage_2_results: CrossFileAnalysisResult, - is_incremental: bool = False - ) -> str: - """ - Stage 3: Generate Markdown report. - In incremental mode, includes summary of resolved vs new issues. - """ - stage_1_json = json.dumps([i.model_dump() for i in stage_1_issues], indent=2) - stage_2_json = stage_2_results.model_dump_json(indent=2) - plan_json = plan.model_dump_json(indent=2) - - # Add incremental context to aggregation - incremental_context = "" - if is_incremental: - resolved_count = sum(1 for i in stage_1_issues if i.isResolved) - new_count = len(stage_1_issues) - resolved_count - previous_count = len(request.previousCodeAnalysisIssues or []) - incremental_context = f""" -## INCREMENTAL REVIEW SUMMARY -- Previous issues from last review: {previous_count} -- Issues resolved in this update: {resolved_count} -- New issues found in delta: {new_count} -- Total issues after reconciliation: {len(stage_1_issues)} -""" - - prompt = PromptBuilder.build_stage_3_aggregation_prompt( - repo_slug=request.projectVcsRepoSlug, - pr_id=str(request.pullRequestId), - author="Unknown", - pr_title=request.prTitle or "", - total_files=len(request.changedFiles or []), - additions=0, # Need accurate stats - deletions=0, - stage_0_plan=plan_json, - stage_1_issues_json=stage_1_json, - stage_2_findings_json=stage_2_json, - recommendation=stage_2_results.pr_recommendation - ) - - response = await self.llm.ainvoke(prompt) - return extract_llm_response_text(response) - - async def _parse_response(self, content: str, model_class: Any, retries: int = 2) -> Any: - """ - Robustly parse JSON response into a Pydantic model with retries. - Falls back to manual parsing if structured output wasn't used. - """ - last_error = None - - # Initial cleaning attempt - try: - cleaned = self._clean_json_text(content) - logger.debug(f"Cleaned JSON for {model_class.__name__} (first 500 chars): {cleaned[:500]}") - data = json.loads(cleaned) - return model_class(**data) - except Exception as e: - last_error = e - logger.warning(f"Initial parse failed for {model_class.__name__}: {e}") - logger.debug(f"Raw content (first 1000 chars): {content[:1000]}") - - # Retry with structured output if available - try: - logger.info(f"Attempting structured output retry for {model_class.__name__}") - structured_llm = self.llm.with_structured_output(model_class) - result = await structured_llm.ainvoke( - f"Parse and return this as valid {model_class.__name__}:\n{content[:4000]}" - ) - if result: - logger.info(f"Structured output retry succeeded for {model_class.__name__}") - return result - except Exception as e: - logger.warning(f"Structured output retry failed: {e}") - last_error = e - - # Final fallback: LLM repair loop - for attempt in range(retries): - try: - logger.info(f"Repairing JSON for {model_class.__name__}, attempt {attempt+1}") - repaired = await self._repair_json_with_llm( - content, - str(last_error), - model_class.model_json_schema() - ) - cleaned = self._clean_json_text(repaired) - logger.debug(f"Repaired JSON attempt {attempt+1} (first 500 chars): {cleaned[:500]}") - data = json.loads(cleaned) - return model_class(**data) - except Exception as e: - last_error = e - logger.warning(f"Retry {attempt+1} failed: {e}") - - raise ValueError(f"Failed to parse {model_class.__name__} after retries: {last_error}") - - async def _repair_json_with_llm(self, broken_json: str, error: str, schema: Any) -> str: - """ - Ask LLM to repair malformed JSON. - """ - # Truncate the broken JSON to avoid token limits but show enough context - truncated_json = broken_json[:3000] if len(broken_json) > 3000 else broken_json - - prompt = f"""You are a JSON repair expert. -The following JSON failed to parse/validate: -Error: {error} - -Broken JSON: -{truncated_json} - -Required Schema (the output MUST be a JSON object, not an array): -{json.dumps(schema, indent=2)} - -CRITICAL INSTRUCTIONS: -1. Return ONLY the fixed valid JSON object -2. The response MUST start with {{ and end with }} -3. All property names MUST be enclosed in double quotes -4. No markdown code blocks (no ```) -5. No explanatory text before or after the JSON -6. Ensure all required fields from the schema are present - -Output the corrected JSON object now:""" - response = await self.llm.ainvoke(prompt) - return extract_llm_response_text(response) - - def _clean_json_text(self, text: str) -> str: - """ - Clean markdown and extraneous text from JSON. - """ - text = text.strip() - - # Remove markdown code blocks - if text.startswith("```"): - lines = text.split("\n") - # Skip the opening ``` line (with or without language identifier) - lines = lines[1:] - # Remove trailing ``` if present - if lines and lines[-1].strip() == "```": - lines = lines[:-1] - text = "\n".join(lines).strip() - - # Also handle case where ``` appears mid-text - if "```json" in text: - start_idx = text.find("```json") - end_idx = text.find("```", start_idx + 7) - if end_idx != -1: - text = text[start_idx + 7:end_idx].strip() - else: - text = text[start_idx + 7:].strip() - elif "```" in text: - # Generic code block without language - start_idx = text.find("```") - remaining = text[start_idx + 3:] - end_idx = remaining.find("```") - if end_idx != -1: - text = remaining[:end_idx].strip() - else: - text = remaining.strip() - - # Find JSON object boundaries - obj_start = text.find("{") - obj_end = text.rfind("}") - arr_start = text.find("[") - arr_end = text.rfind("]") - - # Determine if we have an object or array (whichever comes first) - if obj_start != -1 and obj_end != -1: - if arr_start == -1 or obj_start < arr_start: - # Object comes first or no array - text = text[obj_start:obj_end+1] - elif arr_start < obj_start and arr_end != -1: - # Array comes first - but we need an object for Pydantic - # Check if the object is nested inside the array or separate - if obj_end > arr_end: - # Object extends beyond array - likely the object we want - text = text[obj_start:obj_end+1] - else: - # Try to use the object anyway - text = text[obj_start:obj_end+1] - elif arr_start != -1 and arr_end != -1 and obj_start == -1: - # Only array found - log warning as Pydantic models expect objects - logger.warning(f"JSON cleaning found array instead of object, this may fail parsing") - text = text[arr_start:arr_end+1] - - return text - - def _format_rag_context( - self, - rag_context: Optional[Dict[str, Any]], - relevant_files: Optional[set] = None, - pr_changed_files: Optional[List[str]] = None - ) -> str: - """ - Format RAG context into a readable string for the prompt. - Includes rich AST metadata (imports, extends, implements) for better LLM context. - - Args: - rag_context: RAG response with code chunks - relevant_files: Files in current batch to prioritize - pr_changed_files: ALL files modified in the PR - chunks from these files - are marked as potentially stale (from main branch, not PR branch) - """ - if not rag_context: - logger.debug("RAG context is empty or None") - return "" - - # Handle both "chunks" and "relevant_code" keys (RAG API uses "relevant_code") - chunks = rag_context.get("relevant_code", []) or rag_context.get("chunks", []) - if not chunks: - logger.debug("No chunks found in RAG context (keys: %s)", list(rag_context.keys())) - return "" - - logger.debug(f"Processing {len(chunks)} RAG chunks for context") - - # Normalize PR changed files for comparison - pr_changed_set = set() - if pr_changed_files: - for f in pr_changed_files: - pr_changed_set.add(f) - # Also add just the filename for matching - if "/" in f: - pr_changed_set.add(f.rsplit("/", 1)[-1]) - - formatted_parts = [] - included_count = 0 - skipped_modified = 0 - skipped_relevance = 0 - for chunk in chunks: - if included_count >= 10: # Limit to top 10 chunks for focused context - break - - # Extract metadata - metadata = chunk.get("metadata", {}) - path = metadata.get("path", chunk.get("path", "unknown")) - chunk_type = metadata.get("content_type", metadata.get("type", "code")) - score = chunk.get("score", chunk.get("relevance_score", 0)) - - # Check if this chunk is from a file being modified in the PR - is_from_modified_file = False - if pr_changed_set: - path_filename = path.rsplit("/", 1)[-1] if "/" in path else path - is_from_modified_file = ( - path in pr_changed_set or - path_filename in pr_changed_set or - any(path.endswith(f) or f.endswith(path) for f in pr_changed_set) - ) - - # Skip chunks from files being modified in the PR - they're stale - if is_from_modified_file: - logger.debug(f"Skipping RAG chunk from modified file: {path}") - skipped_modified += 1 - continue - - # Optionally filter by relevance to batch files - if relevant_files: - # Include if the chunk's path relates to any batch file - is_relevant = any( - path in f or f in path or - path.rsplit("/", 1)[-1] == f.rsplit("/", 1)[-1] - for f in relevant_files - ) - # Also include high-scoring chunks regardless - if not is_relevant and score < 0.85: - skipped_relevance += 1 - continue - - text = chunk.get("text", chunk.get("content", "")) - if not text: - continue - - included_count += 1 - - # Build rich metadata context from AST-extracted info - meta_lines = [] - meta_lines.append(f"File: {path}") - - # Include namespace/package if available - if metadata.get("namespace"): - meta_lines.append(f"Namespace: {metadata['namespace']}") - elif metadata.get("package"): - meta_lines.append(f"Package: {metadata['package']}") - - # Include class/function name - if metadata.get("primary_name"): - meta_lines.append(f"Definition: {metadata['primary_name']}") - elif metadata.get("semantic_names"): - meta_lines.append(f"Definitions: {', '.join(metadata['semantic_names'][:5])}") - - # Include inheritance info (extends, implements) - if metadata.get("extends"): - extends = metadata["extends"] - if isinstance(extends, list): - meta_lines.append(f"Extends: {', '.join(extends)}") - else: - meta_lines.append(f"Extends: {extends}") - - if metadata.get("implements"): - implements = metadata["implements"] - if isinstance(implements, list): - meta_lines.append(f"Implements: {', '.join(implements)}") - else: - meta_lines.append(f"Implements: {implements}") - - # Include imports (abbreviated if too many) - if metadata.get("imports"): - imports = metadata["imports"] - if isinstance(imports, list): - if len(imports) <= 5: - meta_lines.append(f"Imports: {'; '.join(imports)}") - else: - meta_lines.append(f"Imports: {'; '.join(imports[:5])}... (+{len(imports)-5} more)") - - # Include parent context (for nested methods) - if metadata.get("parent_context"): - parent_ctx = metadata["parent_context"] - if isinstance(parent_ctx, list): - meta_lines.append(f"Parent: {'.'.join(parent_ctx)}") - - # Include content type for understanding chunk nature - if chunk_type and chunk_type != "code": - meta_lines.append(f"Type: {chunk_type}") - - meta_text = "\n".join(meta_lines) - formatted_parts.append( - f"### Related Code #{included_count} (relevance: {score:.2f})\n" - f"{meta_text}\n" - f"```\n{text}\n```\n" - ) - - if not formatted_parts: - logger.info(f"No RAG chunks included in prompt (total: {len(chunks)}, skipped_modified: {skipped_modified}, skipped_relevance: {skipped_relevance})") - return "" - - logger.info(f"Included {len(formatted_parts)} RAG chunks in prompt context (skipped: {skipped_modified} modified, {skipped_relevance} low relevance)") - return "\n".join(formatted_parts) - - def _emit_status(self, state: str, message: str): - if self.event_callback: - self.event_callback({ - "type": "status", - "state": state, - "message": message - }) - - def _emit_progress(self, percent: int, message: str): - if self.event_callback: - self.event_callback({ - "type": "progress", - "percent": percent, - "message": message - }) - - def _emit_error(self, message: str): - if self.event_callback: - self.event_callback({ - "type": "error", - "message": message - }) - - def _count_files(self, plan: ReviewPlan) -> int: - count = 0 - for group in plan.file_groups: - count += len(group.files) - return count - - def _ensure_all_files_planned(self, plan: ReviewPlan, all_changed_files: List[str]) -> ReviewPlan: - """ - Ensure all changed files are included in the review plan. - If LLM missed files, add them to a LOW priority group. - """ - # Collect files already in the plan - planned_files = set() - for group in plan.file_groups: - for f in group.files: - planned_files.add(f.path) - - # Also count skipped files - skipped_files = set() - for skip in plan.files_to_skip: - skipped_files.add(skip.path) - - # Find missing files - all_files_set = set(all_changed_files) - missing_files = all_files_set - planned_files - skipped_files - - if missing_files: - logger.warning( - f"Stage 0 plan missing {len(missing_files)} files out of {len(all_changed_files)}. " - f"Adding to LOW priority group." - ) - - # Create ReviewFile objects for missing files - missing_review_files = [] - for path in missing_files: - missing_review_files.append(ReviewFile( - path=path, - focus_areas=["GENERAL"], - risk_level="LOW", - estimated_issues=0 - )) - - # Add a new group for missing files or append to existing LOW group - low_group_found = False - for group in plan.file_groups: - if group.priority == "LOW": - group.files.extend(missing_review_files) - low_group_found = True - break - - if not low_group_found: - plan.file_groups.append(FileGroup( - group_id="GROUP_MISSING_FILES", - priority="LOW", - rationale="Files not categorized by planner - added automatically", - files=missing_review_files - )) - - logger.info(f"Plan now includes {self._count_files(plan)} files for review") - else: - logger.info( - f"Stage 0 plan complete: {len(planned_files)} files to review, " - f"{len(skipped_files)} files skipped" - ) - - return plan diff --git a/python-ecosystem/mcp-client/service/pooled_review_service.py b/python-ecosystem/mcp-client/service/pooled_review_service.py deleted file mode 100644 index dce735bb..00000000 --- a/python-ecosystem/mcp-client/service/pooled_review_service.py +++ /dev/null @@ -1,261 +0,0 @@ -""" -Example integration of MCP Process Pool with ReviewService. - -This file shows how to integrate the process pool for improved performance. -The actual integration requires careful testing due to the complexity of -the mcp_use library's MCPAgent/MCPClient. - -IMPORTANT: The mcp_use library creates its own subprocess internally. -To fully benefit from pooling, we need to either: -1. Modify mcp_use to accept an existing process -2. Or bypass mcp_use and communicate directly with our pooled processes - -This example shows approach #2 - direct communication with pooled processes. -""" - -import os -import asyncio -import json -import logging -from typing import Dict, Any, Optional, Callable -from dotenv import load_dotenv - -from utils.mcp_pool import get_mcp_pool, McpProcessPool -from model.models import ReviewRequestDto -from llm.llm_factory import LLMFactory -from utils.prompts.prompt_builder import PromptBuilder -from utils.response_parser import ResponseParser -from service.rag_client import RagClient - -logger = logging.getLogger(__name__) - - -class PooledReviewService: - """ - Review service using process pooling for MCP servers. - - This is an alternative implementation that bypasses mcp_use's internal - process management to use our own process pool. - - Benefits: - - Zero JVM startup overhead after pool warmup - - Shared memory footprint across requests - - Better resource utilization for SaaS deployments - - Trade-offs: - - Requires direct MCP protocol communication - - May not support all mcp_use features - """ - - MAX_FIX_RETRIES = 2 - - def __init__(self): - load_dotenv() - self.default_jar_path = os.environ.get( - "MCP_SERVER_JAR", - "/app/codecrow-vcs-mcp-1.0.jar" - ) - self.rag_client = RagClient() - self._pool: Optional[McpProcessPool] = None - - # Check if pooling is enabled - self.pooling_enabled = os.environ.get("MCP_POOLING_ENABLED", "false").lower() == "true" - - async def _get_pool(self) -> McpProcessPool: - """Get or initialize the process pool.""" - if self._pool is None: - self._pool = await get_mcp_pool(self.default_jar_path) - return self._pool - - async def process_review_request( - self, - request: ReviewRequestDto, - event_callback: Optional[Callable[[Dict], None]] = None - ) -> Dict[str, Any]: - """ - Process a review request. - - If pooling is enabled, uses the process pool. - Otherwise, falls back to the original mcp_use implementation. - """ - if self.pooling_enabled: - return await self._process_review_pooled(request, event_callback) - else: - # Fall back to original implementation - from service.review_service import ReviewService - original_service = ReviewService() - return await original_service.process_review_request(request, event_callback) - - async def _process_review_pooled( - self, - request: ReviewRequestDto, - event_callback: Optional[Callable[[Dict], None]] = None - ) -> Dict[str, Any]: - """ - Process review using pooled MCP server process. - - This method communicates directly with the MCP server via STDIO - instead of using mcp_use's internal process spawning. - """ - try: - self._emit_event(event_callback, { - "type": "status", - "state": "started", - "message": "Analysis starting (pooled mode)" - }) - - pool = await self._get_pool() - - self._emit_event(event_callback, { - "type": "status", - "state": "pool_acquired", - "message": "Acquired MCP server from pool" - }) - - async with pool.acquire() as pooled_process: - # Build the prompt and get LLM - pr_metadata = self._build_pr_metadata(request) - rag_context = await self._fetch_rag_context(request, event_callback) - prompt = self._build_prompt(request, pr_metadata, rag_context) - llm = self._create_llm(request) - - self._emit_event(event_callback, { - "type": "status", - "state": "executing", - "message": "Executing analysis with pooled process" - }) - - # Execute the analysis using the pooled process - # This is a simplified version - full implementation would - # need to handle the MCP protocol properly - result = await self._execute_with_pooled_process( - pooled_process, - prompt, - llm, - request, - event_callback - ) - - self._emit_event(event_callback, { - "type": "final", - "result": "Analysis completed" - }) - - return {"result": result} - - except Exception as e: - logger.error(f"Pooled review failed: {e}", exc_info=True) - error_response = ResponseParser.create_error_response( - "Agent execution failed (pooled)", str(e) - ) - self._emit_event(event_callback, { - "type": "error", - "message": str(e) - }) - return {"result": error_response} - - async def _execute_with_pooled_process( - self, - pooled_process, - prompt: str, - llm, - request: ReviewRequestDto, - event_callback: Optional[Callable[[Dict], None]] - ) -> Dict[str, Any]: - """ - Execute analysis using a pooled MCP server process. - - NOTE: This is a placeholder that shows the concept. - Full implementation requires proper MCP protocol handling. - - The mcp_use library's MCPClient handles the protocol, but it - creates its own subprocess. To use pooling, we'd need to either: - - 1. Fork mcp_use to accept external processes - 2. Implement MCP protocol directly - 3. Use the MCP SDK's Python client directly - - For now, this shows the architecture for option 2/3. - """ - # MCP JSON-RPC message format - # See: https://modelcontextprotocol.io/docs/spec/protocol - - # Initialize connection - init_request = { - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": { - "name": "codecrow-mcp-client", - "version": "1.0.0" - } - } - } - - # Send initialization (this is conceptual - actual impl needs proper I/O handling) - # await self._send_mcp_message(pooled_process, init_request) - # response = await self._read_mcp_response(pooled_process) - - # For now, return a placeholder indicating pooling is working - # but full protocol implementation is pending - return { - "summary": "Pooled execution placeholder", - "issues": [], - "note": "Full MCP protocol integration pending. Process pool is working.", - "pool_metrics": (await self._get_pool()).get_metrics() - } - - def _emit_event(self, callback, event): - """Emit event to callback if provided.""" - if callback: - try: - callback(event) - except Exception as e: - logger.warning(f"Event callback error: {e}") - - def _build_pr_metadata(self, request: ReviewRequestDto) -> str: - """Build PR metadata string.""" - return f"PR #{request.pullRequestId} in {request.workspace}/{request.repoSlug}" - - async def _fetch_rag_context(self, request, event_callback) -> Optional[str]: - """Fetch RAG context if enabled.""" - if not request.ragEnabled: - return None - try: - return await self.rag_client.get_context(request) - except Exception as e: - logger.warning(f"RAG context fetch failed: {e}") - return None - - def _build_prompt(self, request, pr_metadata, rag_context) -> str: - """Build the analysis prompt.""" - return PromptBuilder.build_review_prompt( - pr_metadata=pr_metadata, - rag_context=rag_context, - custom_instructions=request.customInstructions - ) - - def _create_llm(self, request): - """Create LLM instance for the request.""" - return LLMFactory.create( - provider=request.aiProvider, - api_key=request.aiProviderApiKey, - model=request.aiModel, - max_tokens=request.maxAllowedTokens - ) - -#TODO: Implement pooling logic -# Example usage in web server -async def create_pooled_service(): - """Create and initialize a pooled review service.""" - service = PooledReviewService() - - # Pre-warm the pool - if service.pooling_enabled: - pool = await service._get_pool() - logger.info(f"MCP pool initialized: {pool.get_metrics()}") - - return service diff --git a/python-ecosystem/mcp-client/service/rag/__init__.py b/python-ecosystem/mcp-client/service/rag/__init__.py new file mode 100644 index 00000000..b25df464 --- /dev/null +++ b/python-ecosystem/mcp-client/service/rag/__init__.py @@ -0,0 +1,19 @@ +""" +RAG Service Package. + +Contains components for RAG (Retrieval Augmented Generation) functionality: +- rag_client: Client for interacting with the RAG Pipeline API +- llm_reranker: LLM-based reranking for RAG results +""" +from service.rag.rag_client import RagClient, RAG_MIN_RELEVANCE_SCORE, RAG_DEFAULT_TOP_K +from service.rag.llm_reranker import LLMReranker, LLM_RERANK_ENABLED, LLM_RERANK_THRESHOLD, RerankResult + +__all__ = [ + "RagClient", + "RAG_MIN_RELEVANCE_SCORE", + "RAG_DEFAULT_TOP_K", + "LLMReranker", + "LLM_RERANK_ENABLED", + "LLM_RERANK_THRESHOLD", + "RerankResult", +] diff --git a/python-ecosystem/mcp-client/service/llm_reranker.py b/python-ecosystem/mcp-client/service/rag/llm_reranker.py similarity index 100% rename from python-ecosystem/mcp-client/service/llm_reranker.py rename to python-ecosystem/mcp-client/service/rag/llm_reranker.py diff --git a/python-ecosystem/mcp-client/service/rag/rag_client.py b/python-ecosystem/mcp-client/service/rag/rag_client.py new file mode 100644 index 00000000..0c7e4972 --- /dev/null +++ b/python-ecosystem/mcp-client/service/rag/rag_client.py @@ -0,0 +1,405 @@ +""" +RAG Pipeline client for retrieving contextual information during code review. +""" +import os +import logging +from datetime import datetime +from typing import Dict, List, Optional, Any +import httpx + +logger = logging.getLogger(__name__) + +# RAG configuration from environment variables +RAG_MIN_RELEVANCE_SCORE = float(os.environ.get("RAG_MIN_RELEVANCE_SCORE", "0.7")) +RAG_DEFAULT_TOP_K = int(os.environ.get("RAG_DEFAULT_TOP_K", "15")) + + +class RagClient: + """Client for interacting with the RAG Pipeline API.""" + + def __init__(self, base_url: Optional[str] = None, enabled: Optional[bool] = None): + """ + Initialize RAG client. + + Args: + base_url: RAG pipeline API URL (default from env RAG_API_URL) + enabled: Whether RAG is enabled (default from env RAG_ENABLED) + """ + self.base_url = base_url or os.environ.get("RAG_API_URL", "http://rag-pipeline:8001") + self.enabled = enabled if enabled is not None else os.environ.get("RAG_ENABLED", "false").lower() == "true" + self.timeout = 30.0 + self._client: Optional[httpx.AsyncClient] = None + self._service_secret = os.environ.get("SERVICE_SECRET", "") + + if self.enabled: + logger.info(f"RAG client initialized: {self.base_url}") + else: + logger.info("RAG client disabled") + + async def _get_client(self) -> httpx.AsyncClient: + """Get or create an HTTP client for connection pooling (instance-level).""" + if self._client is None or self._client.is_closed: + headers = {} + if self._service_secret: + headers["x-service-secret"] = self._service_secret + self._client = httpx.AsyncClient( + timeout=self.timeout, + limits=httpx.Limits(max_connections=10, max_keepalive_connections=5), + headers=headers, + ) + return self._client + + async def close(self): + """Close this instance's HTTP client.""" + if self._client is not None and not self._client.is_closed: + await self._client.aclose() + self._client = None + + async def get_pr_context( + self, + workspace: str, + project: str, + branch: Optional[str], + changed_files: List[str], + diff_snippets: Optional[List[str]] = None, + pr_title: Optional[str] = None, + pr_description: Optional[str] = None, + top_k: int = None, + enable_priority_reranking: bool = True, + min_relevance_score: float = None, + base_branch: Optional[str] = None, + deleted_files: Optional[List[str]] = None, + pr_number: Optional[int] = None, + all_pr_changed_files: Optional[List[str]] = None + ) -> Dict[str, Any]: + """ + Get relevant context for PR review with multi-branch support. + + Args: + workspace: Workspace identifier + project: Project identifier + branch: Target branch (PR source branch) + changed_files: List of files changed in the PR + diff_snippets: Code snippets extracted from diff for semantic search + pr_title: PR title for semantic understanding + pr_description: Optional PR description + top_k: Number of relevant chunks to retrieve (default from RAG_DEFAULT_TOP_K) + enable_priority_reranking: Enable priority-based score boosting + min_relevance_score: Minimum relevance threshold (default from RAG_MIN_RELEVANCE_SCORE) + base_branch: Base branch (PR target, e.g., 'main'). Auto-detected if not provided. + deleted_files: Files deleted in target branch (excluded from results) + pr_number: If set, enables hybrid query with PR-indexed data priority + all_pr_changed_files: All files in PR (for exclusion from branch query in hybrid mode) + + Returns: + Dict with context information or empty dict if RAG is disabled + """ + if not self.enabled: + logger.debug("RAG disabled, returning empty context") + return {"context": {"relevant_code": []}} + + # Branch is required for RAG query + if not branch: + logger.warning("Branch not specified for RAG query, skipping") + return {"context": {"relevant_code": []}} + + # Apply defaults from env vars + if top_k is None: + top_k = RAG_DEFAULT_TOP_K + if min_relevance_score is None: + min_relevance_score = RAG_MIN_RELEVANCE_SCORE + + start_time = datetime.now() + + try: + payload = { + "workspace": workspace, + "project": project, + "branch": branch, + "changed_files": changed_files, + "diff_snippets": diff_snippets or [], + "pr_title": pr_title, + "pr_description": pr_description, + "top_k": top_k, + "enable_priority_reranking": enable_priority_reranking, + "min_relevance_score": min_relevance_score + } + + # Add optional multi-branch parameters + if base_branch: + payload["base_branch"] = base_branch + if deleted_files: + payload["deleted_files"] = deleted_files + + # Add hybrid mode parameters + if pr_number: + payload["pr_number"] = pr_number + if all_pr_changed_files: + payload["all_pr_changed_files"] = all_pr_changed_files + + client = await self._get_client() + response = await client.post( + f"{self.base_url}/query/pr-context", + json=payload + ) + response.raise_for_status() + result = response.json() + + # Log timing and result stats + elapsed_ms = (datetime.now() - start_time).total_seconds() * 1000 + chunk_count = len(result.get("context", {}).get("relevant_code", [])) + branches_searched = result.get("context", {}).get("_branches_searched", [branch]) + logger.info(f"RAG query completed in {elapsed_ms:.2f}ms, retrieved {chunk_count} chunks from branches: {branches_searched}") + + return result + + except httpx.HTTPError as e: + logger.warning(f"Failed to retrieve PR context from RAG: {e}") + return {"context": {"relevant_code": []}} + except Exception as e: + logger.error(f"Unexpected error querying RAG: {e}") + return {"context": {"relevant_code": []}} + + async def semantic_search( + self, + query: str, + workspace: str, + project: str, + branch: str, + top_k: int = 5, + filter_language: Optional[str] = None + ) -> Dict[str, Any]: + """ + Perform semantic search in the repository. + + Args: + query: Search query + workspace: Workspace identifier + project: Project identifier + branch: Branch name + top_k: Number of results to return + filter_language: Optional language filter (e.g., 'python', 'java') + + Returns: + Dict with search results or empty results if RAG is disabled + """ + if not self.enabled: + return {"results": []} + + try: + payload = { + "query": query, + "workspace": workspace, + "project": project, + "branch": branch, + "top_k": top_k + } + if filter_language: + payload["filter_language"] = filter_language + + client = await self._get_client() + response = await client.post( + f"{self.base_url}/query/search", + json=payload + ) + response.raise_for_status() + return response.json() + + except httpx.HTTPError as e: + logger.warning(f"Semantic search failed: {e}") + return {"results": []} + except Exception as e: + logger.error(f"Unexpected error in semantic search: {e}") + return {"results": []} + + async def is_healthy(self) -> bool: + """ + Check if RAG pipeline is healthy. + + Returns: + True if RAG is enabled and healthy, False otherwise + """ + if not self.enabled: + return False + + try: + client = await self._get_client() + response = await client.get(f"{self.base_url}/health") + return response.status_code == 200 + except Exception as e: + logger.warning(f"RAG health check failed: {e}") + return False + + async def get_deterministic_context( + self, + workspace: str, + project: str, + branches: List[str], + file_paths: List[str], + limit_per_file: int = 10 + ) -> Dict[str, Any]: + """ + Get context using DETERMINISTIC metadata-based retrieval. + + Two-step process leveraging tree-sitter metadata: + 1. Get chunks for the changed file_paths + 2. Extract semantic_names/imports/extends from those chunks + 3. Find related definitions using extracted identifiers + + NO language-specific parsing needed - tree-sitter did it during indexing! + Predictable: same input = same output. + + Args: + workspace: Workspace identifier + project: Project identifier + branches: Branches to search (e.g., ['release/1.29', 'master']) + file_paths: Changed file paths from diff + limit_per_file: Max chunks per file (default 10) + + Returns: + Dict with chunks grouped by: changed_files, related_definitions + """ + if not self.enabled: + logger.debug("RAG disabled, returning empty deterministic context") + return {"context": {"chunks": [], "changed_files": {}, "related_definitions": {}}} + + start_time = datetime.now() + + try: + payload = { + "workspace": workspace, + "project": project, + "branches": branches, + "file_paths": file_paths, + "limit_per_file": limit_per_file + } + + client = await self._get_client() + response = await client.post( + f"{self.base_url}/query/deterministic", + json=payload + ) + response.raise_for_status() + result = response.json() + + # Log timing and stats + elapsed_ms = (datetime.now() - start_time).total_seconds() * 1000 + context = result.get("context", {}) + chunk_count = len(context.get("chunks", [])) + logger.info(f"Deterministic RAG query completed in {elapsed_ms:.2f}ms, " + f"retrieved {chunk_count} chunks for {len(file_paths)} files") + + return result + + except httpx.HTTPError as e: + logger.warning(f"Failed to retrieve deterministic context: {e}") + return {"context": {"chunks": [], "changed_files": {}, "related_definitions": {}}} + except Exception as e: + logger.error(f"Unexpected error in deterministic RAG query: {e}") + return {"context": {"chunks": [], "changed_files": {}, "related_definitions": {}}} + + # ========================================================================= + # PR File Indexing Methods (for PR-specific RAG layer) + # ========================================================================= + + async def index_pr_files( + self, + workspace: str, + project: str, + pr_number: int, + branch: str, + files: List[Dict[str, str]] + ) -> Dict[str, Any]: + """ + Index PR files into the main collection with PR-specific metadata. + + Files are indexed with metadata (pr=true, pr_number=X) to enable + hybrid queries that prioritize PR data over branch data. + + Existing PR points for the same pr_number are deleted first. + + Args: + workspace: Workspace identifier + project: Project identifier + pr_number: PR number for metadata tagging + branch: Source branch name + files: List of {path: str, content: str, change_type: str} + + Returns: + Dict with indexing status and chunk counts + """ + if not self.enabled: + logger.debug("RAG disabled, skipping PR file indexing") + return {"status": "skipped", "chunks_indexed": 0} + + if not files: + logger.debug("No files to index for PR") + return {"status": "skipped", "chunks_indexed": 0} + + try: + payload = { + "workspace": workspace, + "project": project, + "pr_number": pr_number, + "branch": branch, + "files": files + } + + client = await self._get_client() + response = await client.post( + f"{self.base_url}/index/pr-files", + json=payload, + timeout=120.0 # Longer timeout for indexing + ) + response.raise_for_status() + result = response.json() + + logger.info(f"Indexed PR #{pr_number}: {result.get('chunks_indexed', 0)} chunks from {result.get('files_processed', 0)} files") + return result + + except httpx.HTTPError as e: + logger.warning(f"Failed to index PR files: {e}") + return {"status": "error", "error": str(e)} + except Exception as e: + logger.error(f"Unexpected error indexing PR files: {e}") + return {"status": "error", "error": str(e)} + + async def delete_pr_files( + self, + workspace: str, + project: str, + pr_number: int + ) -> bool: + """ + Delete all indexed points for a specific PR. + + Called after analysis completes to clean up PR-specific data. + + Args: + workspace: Workspace identifier + project: Project identifier + pr_number: PR number to delete + + Returns: + True if deleted successfully, False otherwise + """ + if not self.enabled: + return True + + try: + client = await self._get_client() + response = await client.delete( + f"{self.base_url}/index/pr-files/{workspace}/{project}/{pr_number}" + ) + response.raise_for_status() + result = response.json() + + logger.info(f"Deleted PR #{pr_number} indexed data") + return result.get("status") == "deleted" + + except httpx.HTTPError as e: + logger.warning(f"Failed to delete PR files: {e}") + return False + except Exception as e: + logger.error(f"Unexpected error deleting PR files: {e}") + return False diff --git a/python-ecosystem/mcp-client/service/rag_client.py b/python-ecosystem/mcp-client/service/rag_client.py deleted file mode 100644 index 44c98d34..00000000 --- a/python-ecosystem/mcp-client/service/rag_client.py +++ /dev/null @@ -1,219 +0,0 @@ -""" -RAG Pipeline client for retrieving contextual information during code review. -""" -import os -import logging -from datetime import datetime -from typing import Dict, List, Optional, Any -import httpx - -logger = logging.getLogger(__name__) - -# RAG configuration from environment variables -RAG_MIN_RELEVANCE_SCORE = float(os.environ.get("RAG_MIN_RELEVANCE_SCORE", "0.7")) -RAG_DEFAULT_TOP_K = int(os.environ.get("RAG_DEFAULT_TOP_K", "15")) - - -class RagClient: - """Client for interacting with the RAG Pipeline API.""" - - # Shared HTTP client for connection pooling - _shared_client: Optional[httpx.AsyncClient] = None - - def __init__(self, base_url: Optional[str] = None, enabled: Optional[bool] = None): - """ - Initialize RAG client. - - Args: - base_url: RAG pipeline API URL (default from env RAG_API_URL) - enabled: Whether RAG is enabled (default from env RAG_ENABLED) - """ - self.base_url = base_url or os.environ.get("RAG_API_URL", "http://rag-pipeline:8001") - self.enabled = enabled if enabled is not None else os.environ.get("RAG_ENABLED", "false").lower() == "true" - self.timeout = 30.0 - - if self.enabled: - logger.info(f"RAG client initialized: {self.base_url}") - else: - logger.info("RAG client disabled") - - async def _get_client(self) -> httpx.AsyncClient: - """Get or create a shared HTTP client for connection pooling.""" - if RagClient._shared_client is None or RagClient._shared_client.is_closed: - RagClient._shared_client = httpx.AsyncClient( - timeout=self.timeout, - limits=httpx.Limits(max_connections=10, max_keepalive_connections=5) - ) - return RagClient._shared_client - - async def close(self): - """Close the shared HTTP client.""" - if RagClient._shared_client is not None and not RagClient._shared_client.is_closed: - await RagClient._shared_client.aclose() - RagClient._shared_client = None - - async def get_pr_context( - self, - workspace: str, - project: str, - branch: Optional[str], - changed_files: List[str], - diff_snippets: Optional[List[str]] = None, - pr_title: Optional[str] = None, - pr_description: Optional[str] = None, - top_k: int = None, - enable_priority_reranking: bool = True, - min_relevance_score: float = None, - base_branch: Optional[str] = None, - deleted_files: Optional[List[str]] = None - ) -> Dict[str, Any]: - """ - Get relevant context for PR review with multi-branch support. - - Args: - workspace: Workspace identifier - project: Project identifier - branch: Target branch (PR source branch) - changed_files: List of files changed in the PR - diff_snippets: Code snippets extracted from diff for semantic search - pr_title: PR title for semantic understanding - pr_description: Optional PR description - top_k: Number of relevant chunks to retrieve (default from RAG_DEFAULT_TOP_K) - enable_priority_reranking: Enable priority-based score boosting - min_relevance_score: Minimum relevance threshold (default from RAG_MIN_RELEVANCE_SCORE) - base_branch: Base branch (PR target, e.g., 'main'). Auto-detected if not provided. - deleted_files: Files deleted in target branch (excluded from results) - - Returns: - Dict with context information or empty dict if RAG is disabled - """ - if not self.enabled: - logger.debug("RAG disabled, returning empty context") - return {"context": {"relevant_code": []}} - - # Branch is required for RAG query - if not branch: - logger.warning("Branch not specified for RAG query, skipping") - return {"context": {"relevant_code": []}} - - # Apply defaults from env vars - if top_k is None: - top_k = RAG_DEFAULT_TOP_K - if min_relevance_score is None: - min_relevance_score = RAG_MIN_RELEVANCE_SCORE - - start_time = datetime.now() - - try: - payload = { - "workspace": workspace, - "project": project, - "branch": branch, - "changed_files": changed_files, - "diff_snippets": diff_snippets or [], - "pr_title": pr_title, - "pr_description": pr_description, - "top_k": top_k, - "enable_priority_reranking": enable_priority_reranking, - "min_relevance_score": min_relevance_score - } - - # Add optional multi-branch parameters - if base_branch: - payload["base_branch"] = base_branch - if deleted_files: - payload["deleted_files"] = deleted_files - - client = await self._get_client() - response = await client.post( - f"{self.base_url}/query/pr-context", - json=payload - ) - response.raise_for_status() - result = response.json() - - # Log timing and result stats - elapsed_ms = (datetime.now() - start_time).total_seconds() * 1000 - chunk_count = len(result.get("context", {}).get("relevant_code", [])) - branches_searched = result.get("context", {}).get("_branches_searched", [branch]) - logger.info(f"RAG query completed in {elapsed_ms:.2f}ms, retrieved {chunk_count} chunks from branches: {branches_searched}") - - return result - - except httpx.HTTPError as e: - logger.warning(f"Failed to retrieve PR context from RAG: {e}") - return {"context": {"relevant_code": []}} - except Exception as e: - logger.error(f"Unexpected error querying RAG: {e}") - return {"context": {"relevant_code": []}} - - async def semantic_search( - self, - query: str, - workspace: str, - project: str, - branch: str, - top_k: int = 5, - filter_language: Optional[str] = None - ) -> Dict[str, Any]: - """ - Perform semantic search in the repository. - - Args: - query: Search query - workspace: Workspace identifier - project: Project identifier - branch: Branch name - top_k: Number of results to return - filter_language: Optional language filter (e.g., 'python', 'java') - - Returns: - Dict with search results or empty results if RAG is disabled - """ - if not self.enabled: - return {"results": []} - - try: - payload = { - "query": query, - "workspace": workspace, - "project": project, - "branch": branch, - "top_k": top_k - } - if filter_language: - payload["filter_language"] = filter_language - - client = await self._get_client() - response = await client.post( - f"{self.base_url}/query/search", - json=payload - ) - response.raise_for_status() - return response.json() - - except httpx.HTTPError as e: - logger.warning(f"Semantic search failed: {e}") - return {"results": []} - except Exception as e: - logger.error(f"Unexpected error in semantic search: {e}") - return {"results": []} - - async def is_healthy(self) -> bool: - """ - Check if RAG pipeline is healthy. - - Returns: - True if RAG is enabled and healthy, False otherwise - """ - if not self.enabled: - return False - - try: - client = await self._get_client() - response = await client.get(f"{self.base_url}/health") - return response.status_code == 200 - except Exception as e: - logger.warning(f"RAG health check failed: {e}") - return False - diff --git a/python-ecosystem/mcp-client/service/review/__init__.py b/python-ecosystem/mcp-client/service/review/__init__.py new file mode 100644 index 00000000..d136b174 --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/__init__.py @@ -0,0 +1,39 @@ +""" +Review Service Package. + +Contains components for code review functionality: +- review_service: Main entry point for review requests +- orchestrator/: Multi-stage review orchestrator (subpackage) +- issue_processor: Issue post-processing and deduplication +""" +from service.review.review_service import ReviewService +from service.review.orchestrator import ( + MultiStageReviewOrchestrator, + RecursiveMCPAgent, + extract_llm_response_text, + parse_llm_response, + clean_json_text, + reconcile_previous_issues, + deduplicate_issues, +) +from service.review.issue_processor import ( + IssuePostProcessor, + IssueDeduplicator, + post_process_analysis_result, + restore_missing_diffs_from_previous, +) + +__all__ = [ + "ReviewService", + "MultiStageReviewOrchestrator", + "IssuePostProcessor", + "IssueDeduplicator", + "post_process_analysis_result", + "restore_missing_diffs_from_previous", + "RecursiveMCPAgent", + "extract_llm_response_text", + "parse_llm_response", + "clean_json_text", + "reconcile_previous_issues", + "deduplicate_issues", +] diff --git a/python-ecosystem/mcp-client/service/issue_post_processor.py b/python-ecosystem/mcp-client/service/review/issue_processor.py similarity index 100% rename from python-ecosystem/mcp-client/service/issue_post_processor.py rename to python-ecosystem/mcp-client/service/review/issue_processor.py diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/__init__.py b/python-ecosystem/mcp-client/service/review/orchestrator/__init__.py new file mode 100644 index 00000000..36f086a5 --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/__init__.py @@ -0,0 +1,25 @@ +""" +Orchestrator Package. + +Multi-stage review orchestrator components: +- orchestrator: Main MultiStageReviewOrchestrator class +- stages: Stage 0-3 execution methods +- reconciliation: Issue reconciliation and deduplication +- context_helpers: RAG context and diff extraction +- json_utils: JSON parsing and repair +- agents: MCP agent with recursion limit support +""" +from service.review.orchestrator.orchestrator import MultiStageReviewOrchestrator +from service.review.orchestrator.agents import RecursiveMCPAgent, extract_llm_response_text +from service.review.orchestrator.json_utils import parse_llm_response, clean_json_text +from service.review.orchestrator.reconciliation import reconcile_previous_issues, deduplicate_issues + +__all__ = [ + "MultiStageReviewOrchestrator", + "RecursiveMCPAgent", + "extract_llm_response_text", + "parse_llm_response", + "clean_json_text", + "reconcile_previous_issues", + "deduplicate_issues", +] diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/agents.py b/python-ecosystem/mcp-client/service/review/orchestrator/agents.py new file mode 100644 index 00000000..859ce6fc --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/agents.py @@ -0,0 +1,81 @@ +""" +Custom MCP Agent with increased recursion limit. +""" +import logging +from typing import Any +from mcp_use import MCPAgent + +logger = logging.getLogger(__name__) + + +def extract_llm_response_text(response: Any) -> str: + """ + Extract text content from LLM response, handling different response formats. + Some LLM providers return content as a list of objects instead of a string. + """ + if hasattr(response, 'content'): + content = response.content + if isinstance(content, list): + # Handle list content (e.g., from Gemini or other providers) + text_parts = [] + for item in content: + if isinstance(item, str): + text_parts.append(item) + elif isinstance(item, dict): + if 'text' in item: + text_parts.append(item['text']) + elif 'content' in item: + text_parts.append(item['content']) + elif hasattr(item, 'text'): + text_parts.append(item.text) + return "".join(text_parts) + return str(content) + return str(response) + + +# Prevent duplicate logs from mcp_use +mcp_logger = logging.getLogger("mcp_use") +mcp_logger.propagate = False +if not mcp_logger.handlers: + handler = logging.StreamHandler() + formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + handler.setFormatter(formatter) + mcp_logger.addHandler(handler) + + +class RecursiveMCPAgent(MCPAgent): + """ + Subclass of MCPAgent that enforces a higher recursion limit on the internal agent executor. + """ + def __init__(self, *args, recursion_limit: int = 50, **kwargs): + self._custom_recursion_limit = recursion_limit + super().__init__(*args, **kwargs) + + async def stream(self, *args, **kwargs): + """ + Override stream to ensure recursion_limit is applied. + """ + # Ensure the executor exists + if self._agent_executor is None: + await self.initialize() + + # Patch the executor's astream if not already patched + executor = self._agent_executor + if executor and not getattr(executor, "_is_patched_recursion", False): + original_astream = executor.astream + limit = self._custom_recursion_limit + + async def patched_astream(input_data, config=None, **astream_kwargs): + if config is None: + config = {} + config["recursion_limit"] = limit + async for chunk in original_astream(input_data, config=config, **astream_kwargs): + yield chunk + + executor.astream = patched_astream + executor._is_patched_recursion = True + logger.info(f"RecursiveMCPAgent: Patched recursion limit to {limit}") + + # Call parent stream + async for item in super().stream(*args, **kwargs): + yield item diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/context_helpers.py b/python-ecosystem/mcp-client/service/review/orchestrator/context_helpers.py new file mode 100644 index 00000000..b70c3a23 --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/context_helpers.py @@ -0,0 +1,245 @@ +""" +Context and diff extraction helpers for the review orchestrator. +""" +import re +import logging +from typing import Any, Dict, List, Optional, Set + +logger = logging.getLogger(__name__) + + +def extract_symbols_from_diff(diff_content: str) -> List[str]: + """ + Extract potential symbols (identifiers, class names, function names) from diff. + Used to query cross-file context for related changes. + """ + if not diff_content: + return [] + + # Common language keywords/stop-words to filter out + STOP_WORDS = { + # Python + 'import', 'from', 'class', 'def', 'return', 'if', 'else', 'elif', + 'for', 'while', 'try', 'except', 'finally', 'with', 'as', 'pass', + 'break', 'continue', 'raise', 'yield', 'lambda', 'async', 'await', + 'True', 'False', 'None', 'and', 'or', 'not', 'in', 'is', + # Java/TS/JS + 'public', 'private', 'protected', 'static', 'final', 'void', + 'new', 'this', 'super', 'extends', 'implements', 'interface', + 'abstract', 'const', 'let', 'var', 'function', 'export', 'default', + 'throw', 'throws', 'catch', 'instanceof', 'typeof', 'null', + # Common + 'true', 'false', 'null', 'undefined', 'self', 'args', 'kwargs', + 'string', 'number', 'boolean', 'object', 'array', 'list', 'dict', + } + + symbols = set() + + # Patterns for common identifiers + # Match CamelCase identifiers (likely class/component names) + camel_case = re.findall(r'\b([A-Z][a-z]+[A-Z][a-zA-Z]*)\b', diff_content) + symbols.update(camel_case) + + # Match snake_case identifiers (variables, functions) + snake_case = re.findall(r'\b([a-z][a-z0-9]*(?:_[a-z0-9]+)+)\b', diff_content) + symbols.update(s for s in snake_case if len(s) > 5) # Filter short ones + + # Match assignments and function calls + assignments = re.findall(r'\b(\w+)\s*[=:]\s*', diff_content) + symbols.update(a for a in assignments if len(a) > 3) + + # Match import statements + imports = re.findall(r'(?:from|import)\s+([a-zA-Z_][a-zA-Z0-9_.]+)', diff_content) + symbols.update(imports) + + # Filter out stop-words and return + filtered = [s for s in symbols if s.lower() not in STOP_WORDS and len(s) > 2] + return filtered[:20] # Limit to top 20 symbols + + +def extract_diff_snippets(diff_content: str) -> List[str]: + """ + Extract meaningful code snippets from diff content for RAG semantic search. + Focuses on added/modified lines that represent significant code changes. + """ + if not diff_content: + return [] + + snippets = [] + current_snippet_lines = [] + + for line in diff_content.splitlines(): + # Focus on added lines (new code) + if line.startswith("+") and not line.startswith("+++"): + clean_line = line[1:].strip() + # Skip trivial lines + if (clean_line and + len(clean_line) > 10 and # Minimum meaningful length + not clean_line.startswith("//") and # Skip comments + not clean_line.startswith("#") and + not clean_line.startswith("*") and + not clean_line == "{" and + not clean_line == "}" and + not clean_line == ""): + current_snippet_lines.append(clean_line) + + # Batch into snippets of 3-5 lines + if len(current_snippet_lines) >= 3: + snippets.append(" ".join(current_snippet_lines)) + current_snippet_lines = [] + + # Add remaining lines as final snippet + if current_snippet_lines: + snippets.append(" ".join(current_snippet_lines)) + + # Limit to most significant snippets + return snippets[:10] + + +def get_diff_snippets_for_batch( + all_diff_snippets: List[str], + batch_file_paths: List[str] +) -> List[str]: + """ + Filter diff snippets to only include those relevant to the batch files. + + Note: Java DiffParser.extractDiffSnippets() returns CLEAN CODE SNIPPETS (no file paths). + These snippets are just significant code lines like function signatures. + Since snippets don't contain file paths, we return all snippets for semantic search. + The embedding similarity will naturally prioritize relevant matches. + """ + if not all_diff_snippets: + return [] + + # Java snippets are clean code (no file paths), so we can't filter by path + # Return all snippets - the semantic search will find relevant matches + logger.info(f"Using {len(all_diff_snippets)} diff snippets for batch files {batch_file_paths}") + return all_diff_snippets + + +def format_rag_context( + rag_context: Optional[Dict[str, Any]], + relevant_files: Optional[Set[str]] = None, + pr_changed_files: Optional[List[str]] = None +) -> str: + """ + Format RAG context into a readable string for the prompt. + + IMPORTANT: We trust RAG's semantic similarity scores for relevance. + The RAG system already uses embeddings to find semantically related code. + We only filter out chunks from files being modified in the PR (stale data from main branch). + + Args: + rag_context: RAG response with code chunks + relevant_files: (UNUSED - kept for API compatibility) - we trust RAG scores instead + pr_changed_files: Files modified in the PR - chunks from these may be stale + """ + if not rag_context: + logger.debug("RAG context is empty or None") + return "" + + # Handle both "chunks" and "relevant_code" keys (RAG API uses "relevant_code") + chunks = rag_context.get("relevant_code", []) or rag_context.get("chunks", []) + if not chunks: + logger.debug("No chunks found in RAG context (keys: %s)", list(rag_context.keys())) + return "" + + logger.info(f"Processing {len(chunks)} RAG chunks (trusting semantic similarity scores)") + + # Normalize PR changed files for stale-data detection only + pr_changed_set = set() + if pr_changed_files: + for f in pr_changed_files: + pr_changed_set.add(f) + if "/" in f: + pr_changed_set.add(f.rsplit("/", 1)[-1]) + + formatted_parts = [] + included_count = 0 + skipped_stale = 0 + + for chunk in chunks: + if included_count >= 15: + logger.debug(f"Reached chunk limit of 15") + break + + metadata = chunk.get("metadata", {}) + # Support both 'path' and 'file_path' keys (deterministic uses file_path) + path = metadata.get("path") or chunk.get("path") or chunk.get("file_path", "unknown") + chunk_type = metadata.get("content_type", metadata.get("type", "code")) + score = chunk.get("score", chunk.get("relevance_score", 0)) + + # Only filter: chunks from PR-modified files with LOW scores (likely stale) + # High-score chunks from modified files may still be relevant (other parts of same file) + if pr_changed_set: + path_filename = path.rsplit("/", 1)[-1] if "/" in path else path + is_from_modified_file = ( + path in pr_changed_set or + path_filename in pr_changed_set or + any(path.endswith(f) or f.endswith(path) for f in pr_changed_set) + ) + + # Skip ONLY low-score chunks from modified files (likely stale/outdated) + if is_from_modified_file and score < 0.70: + logger.debug(f"Skipping stale chunk from modified file: {path} (score={score:.2f})") + skipped_stale += 1 + continue + + text = chunk.get("text", chunk.get("content", "")) + if not text: + continue + + included_count += 1 + + # Build rich metadata context + meta_lines = [f"File: {path}"] + + if metadata.get("namespace"): + meta_lines.append(f"Namespace: {metadata['namespace']}") + elif metadata.get("package"): + meta_lines.append(f"Package: {metadata['package']}") + + if metadata.get("primary_name"): + meta_lines.append(f"Definition: {metadata['primary_name']}") + elif metadata.get("semantic_names"): + meta_lines.append(f"Definitions: {', '.join(metadata['semantic_names'][:5])}") + + if metadata.get("extends"): + extends = metadata["extends"] + meta_lines.append(f"Extends: {', '.join(extends) if isinstance(extends, list) else extends}") + + if metadata.get("implements"): + implements = metadata["implements"] + meta_lines.append(f"Implements: {', '.join(implements) if isinstance(implements, list) else implements}") + + if metadata.get("imports"): + imports = metadata["imports"] + if isinstance(imports, list): + if len(imports) <= 5: + meta_lines.append(f"Imports: {'; '.join(imports)}") + else: + meta_lines.append(f"Imports: {'; '.join(imports[:5])}... (+{len(imports)-5} more)") + + if metadata.get("parent_context"): + parent_ctx = metadata["parent_context"] + if isinstance(parent_ctx, list): + meta_lines.append(f"Parent: {'.'.join(parent_ctx)}") + + if chunk_type and chunk_type != "code": + meta_lines.append(f"Type: {chunk_type}") + + meta_text = "\n".join(meta_lines) + # Use file path as primary identifier, not a number + # This encourages AI to reference by path rather than by chunk number + formatted_parts.append( + f"### Context from `{path}` (relevance: {score:.2f})\n" + f"{meta_text}\n" + f"```\n{text}\n```\n" + ) + + if not formatted_parts: + logger.warning(f"No RAG chunks included (total: {len(chunks)}, skipped_stale: {skipped_stale})") + return "" + + logger.info(f"Included {len(formatted_parts)} RAG chunks (skipped {skipped_stale} stale from modified files)") + return "\n".join(formatted_parts) diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/json_utils.py b/python-ecosystem/mcp-client/service/review/orchestrator/json_utils.py new file mode 100644 index 00000000..82e15b03 --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/json_utils.py @@ -0,0 +1,155 @@ +""" +JSON parsing, repair, and cleaning utilities for LLM responses. +""" +import json +import logging +from typing import Any, Dict, Optional + +from service.review.orchestrator.agents import extract_llm_response_text + +logger = logging.getLogger(__name__) + + +async def parse_llm_response(content: str, model_class: Any, llm, retries: int = 2) -> Any: + """ + Robustly parse JSON response into a Pydantic model with retries. + Falls back to manual parsing if structured output wasn't used. + """ + last_error = None + + # Initial cleaning attempt + try: + cleaned = clean_json_text(content) + logger.debug(f"Cleaned JSON for {model_class.__name__} (first 500 chars): {cleaned[:500]}") + data = json.loads(cleaned) + return model_class(**data) + except Exception as e: + last_error = e + logger.warning(f"Initial parse failed for {model_class.__name__}: {e}") + logger.debug(f"Raw content (first 1000 chars): {content[:1000]}") + + # Retry with structured output if available + try: + logger.info(f"Attempting structured output retry for {model_class.__name__}") + structured_llm = llm.with_structured_output(model_class) + result = await structured_llm.ainvoke( + f"Parse and return this as valid {model_class.__name__}:\n{content[:4000]}" + ) + if result: + logger.info(f"Structured output retry succeeded for {model_class.__name__}") + return result + except Exception as e: + logger.warning(f"Structured output retry failed: {e}") + last_error = e + + # Final fallback: LLM repair loop + for attempt in range(retries): + try: + logger.info(f"Repairing JSON for {model_class.__name__}, attempt {attempt+1}") + repaired = await repair_json_with_llm( + llm, + content, + str(last_error), + model_class.model_json_schema() + ) + cleaned = clean_json_text(repaired) + logger.debug(f"Repaired JSON attempt {attempt+1} (first 500 chars): {cleaned[:500]}") + data = json.loads(cleaned) + return model_class(**data) + except Exception as e: + last_error = e + logger.warning(f"Retry {attempt+1} failed: {e}") + + raise ValueError(f"Failed to parse {model_class.__name__} after retries: {last_error}") + + +async def repair_json_with_llm(llm, broken_json: str, error: str, schema: Any) -> str: + """ + Ask LLM to repair malformed JSON. + """ + # Truncate the broken JSON to avoid token limits but show enough context + truncated_json = broken_json[:3000] if len(broken_json) > 3000 else broken_json + + prompt = f"""You are a JSON repair expert. +The following JSON failed to parse/validate: +Error: {error} + +Broken JSON: +{truncated_json} + +Required Schema (the output MUST be a JSON object, not an array): +{json.dumps(schema, indent=2)} + +CRITICAL INSTRUCTIONS: +1. Return ONLY the fixed valid JSON object +2. The response MUST start with {{ and end with }} +3. All property names MUST be enclosed in double quotes +4. No markdown code blocks (no ```) +5. No explanatory text before or after the JSON +6. Ensure all required fields from the schema are present + +Output the corrected JSON object now:""" + response = await llm.ainvoke(prompt) + return extract_llm_response_text(response) + + +def clean_json_text(text: str) -> str: + """ + Clean markdown and extraneous text from JSON. + """ + text = text.strip() + + # Remove markdown code blocks + if text.startswith("```"): + lines = text.split("\n") + # Skip the opening ``` line (with or without language identifier) + lines = lines[1:] + # Remove trailing ``` if present + if lines and lines[-1].strip() == "```": + lines = lines[:-1] + text = "\n".join(lines).strip() + + # Also handle case where ``` appears mid-text + if "```json" in text: + start_idx = text.find("```json") + end_idx = text.find("```", start_idx + 7) + if end_idx != -1: + text = text[start_idx + 7:end_idx].strip() + else: + text = text[start_idx + 7:].strip() + elif "```" in text: + # Generic code block without language + start_idx = text.find("```") + remaining = text[start_idx + 3:] + end_idx = remaining.find("```") + if end_idx != -1: + text = remaining[:end_idx].strip() + else: + text = remaining.strip() + + # Find JSON object boundaries + obj_start = text.find("{") + obj_end = text.rfind("}") + arr_start = text.find("[") + arr_end = text.rfind("]") + + # Determine if we have an object or array (whichever comes first) + if obj_start != -1 and obj_end != -1: + if arr_start == -1 or obj_start < arr_start: + # Object comes first or no array + text = text[obj_start:obj_end+1] + elif arr_start < obj_start and arr_end != -1: + # Array comes first - but we need an object for Pydantic + # Check if the object is nested inside the array or separate + if obj_end > arr_end: + # Object extends beyond array - likely the object we want + text = text[obj_start:obj_end+1] + else: + # Try to use the object anyway + text = text[obj_start:obj_end+1] + elif arr_start != -1 and arr_end != -1 and obj_start == -1: + # Only array found - log warning as Pydantic models expect objects + logger.warning(f"JSON cleaning found array instead of object, this may fail parsing") + text = text[arr_start:arr_end+1] + + return text diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/orchestrator.py b/python-ecosystem/mcp-client/service/review/orchestrator/orchestrator.py new file mode 100644 index 00000000..f6f138bd --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/orchestrator.py @@ -0,0 +1,251 @@ +""" +Multi-Stage Review Orchestrator. + +Orchestrates the 4-stage AI code review pipeline: +- Stage 0: Planning & Prioritization +- Stage 1: Parallel File Review +- Stage 2: Cross-File & Architectural Analysis +- Stage 3: Aggregation & Final Report +""" +import logging +from typing import Dict, Any, List, Optional, Callable + +from model.dtos import ReviewRequestDto +from utils.diff_processor import ProcessedDiff + +from service.review.orchestrator.reconciliation import reconcile_previous_issues +from service.review.orchestrator.stages import ( + execute_branch_analysis, + execute_stage_0_planning, + execute_stage_1_file_reviews, + execute_stage_2_cross_file, + execute_stage_3_aggregation, + _emit_status, + _emit_progress, + _emit_error, +) + +logger = logging.getLogger(__name__) + + +class MultiStageReviewOrchestrator: + """ + Orchestrates the 4-stage AI code review pipeline: + Stage 0: Planning & Prioritization + Stage 1: Parallel File Review + Stage 2: Cross-File & Architectural Analysis + Stage 3: Aggregation & Final Report + """ + + def __init__( + self, + llm, + mcp_client, + rag_client=None, + event_callback: Optional[Callable[[Dict], None]] = None + ): + self.llm = llm + self.client = mcp_client + self.rag_client = rag_client + self.event_callback = event_callback + self.max_parallel_stage_1 = 5 + self._pr_number: Optional[int] = None + self._pr_indexed: bool = False + + async def _index_pr_files( + self, + request: ReviewRequestDto, + processed_diff: Optional[ProcessedDiff] + ) -> None: + """ + Index PR files into the main RAG collection with PR-specific metadata. + This enables hybrid queries that prioritize PR data over stale branch data. + """ + if not self.rag_client or not processed_diff: + return + + pr_number = request.pullRequestId + if not pr_number: + logger.info("No PR number, skipping PR file indexing") + return + + files = [] + for f in processed_diff.get_included_files(): + content = f.full_content or f.content + change_type = f.change_type.value if hasattr(f.change_type, 'value') else str(f.change_type) + if content and change_type != "DELETED": + files.append({ + "path": f.path, + "content": content, + "change_type": change_type + }) + + if not files: + logger.info("No files to index for PR") + return + + try: + result = await self.rag_client.index_pr_files( + workspace=request.projectWorkspace, + project=request.projectNamespace, + pr_number=pr_number, + branch=request.targetBranchName or "unknown", + files=files + ) + if result.get("status") == "indexed": + self._pr_number = pr_number + self._pr_indexed = True + logger.info(f"Indexed PR #{pr_number}: {result.get('chunks_indexed', 0)} chunks") + else: + logger.warning(f"Failed to index PR files: {result}") + except Exception as e: + logger.warning(f"Error indexing PR files: {e}") + + async def _cleanup_pr_files(self, request: ReviewRequestDto) -> None: + """Delete PR-indexed data after analysis completes.""" + if not self._pr_indexed or not self._pr_number or not self.rag_client: + return + + try: + await self.rag_client.delete_pr_files( + workspace=request.projectWorkspace, + project=request.projectNamespace, + pr_number=self._pr_number + ) + logger.info(f"Cleaned up PR #{self._pr_number} indexed data") + except Exception as e: + logger.warning(f"Failed to cleanup PR files: {e}") + finally: + self._pr_number = None + self._pr_indexed = False + + async def execute_branch_analysis(self, prompt: str) -> Dict[str, Any]: + """ + Execute a single-pass branch analysis using the provided prompt. + """ + return await execute_branch_analysis( + self.llm, + self.client, + prompt, + self.event_callback + ) + + async def orchestrate_review( + self, + request: ReviewRequestDto, + rag_context: Optional[Dict[str, Any]] = None, + processed_diff: Optional[ProcessedDiff] = None + ) -> Dict[str, Any]: + """ + Main entry point for the multi-stage review. + Supports both FULL (initial review) and INCREMENTAL (follow-up review) modes. + """ + is_incremental = ( + request.analysisMode == "INCREMENTAL" + and request.deltaDiff + ) + + if is_incremental: + logger.info(f"INCREMENTAL mode: reviewing delta diff, {len(request.previousCodeAnalysisIssues or [])} previous issues to reconcile") + else: + logger.info("FULL mode: initial PR review") + + try: + # Index PR files into RAG for hybrid queries + await self._index_pr_files(request, processed_diff) + + # === STAGE 0: Planning === + _emit_status(self.event_callback, "stage_0_started", "Stage 0: Planning & Prioritization...") + review_plan = await execute_stage_0_planning( + self.llm, request, is_incremental, processed_diff=processed_diff + ) + + review_plan = self._ensure_all_files_planned(review_plan, request.changedFiles or []) + _emit_progress(self.event_callback, 10, "Stage 0 Complete: Review plan created") + + # === STAGE 1: File Reviews === + _emit_status(self.event_callback, "stage_1_started", f"Stage 1: Analyzing {self._count_files(review_plan)} files...") + file_issues = await execute_stage_1_file_reviews( + self.llm, + request, + review_plan, + self.rag_client, + rag_context, + processed_diff, + is_incremental, + self.max_parallel_stage_1, + self.event_callback, + self._pr_indexed + ) + _emit_progress(self.event_callback, 60, f"Stage 1 Complete: {len(file_issues)} issues found across files") + + # === STAGE 1.5: Issue Reconciliation === + if request.previousCodeAnalysisIssues: + _emit_status(self.event_callback, "reconciliation_started", "Reconciling previous issues...") + file_issues = await reconcile_previous_issues( + request, file_issues, processed_diff + ) + _emit_progress(self.event_callback, 70, f"Reconciliation Complete: {len(file_issues)} total issues after reconciliation") + + # === STAGE 2: Cross-File Analysis === + _emit_status(self.event_callback, "stage_2_started", "Stage 2: Analyzing cross-file patterns...") + cross_file_results = await execute_stage_2_cross_file( + self.llm, request, file_issues, review_plan, + processed_diff=processed_diff, + ) + _emit_progress(self.event_callback, 85, "Stage 2 Complete: Cross-file analysis finished") + + # === STAGE 3: Aggregation === + _emit_status(self.event_callback, "stage_3_started", "Stage 3: Generating final report...") + final_report = await execute_stage_3_aggregation( + self.llm, request, review_plan, file_issues, cross_file_results, + is_incremental, processed_diff=processed_diff + ) + _emit_progress(self.event_callback, 100, "Stage 3 Complete: Report generated") + + return { + "comment": final_report, + "issues": [issue.model_dump() for issue in file_issues], + } + + except Exception as e: + logger.error(f"Multi-stage review failed: {e}", exc_info=True) + _emit_error(self.event_callback, str(e)) + raise + finally: + await self._cleanup_pr_files(request) + + def _count_files(self, plan) -> int: + """Count total files in review plan.""" + return sum(len(g.files) for g in plan.file_groups) + + def _ensure_all_files_planned(self, plan, changed_files: List[str]): + """ + Ensure all changed files are included in the review plan. + LLM may miss some files, so we add them to a catch-all group. + """ + from model.multi_stage import ReviewFile, FileGroup + + planned_files = set() + for group in plan.file_groups: + for f in group.files: + planned_files.add(f.path) + + missing_files = [f for f in changed_files if f not in planned_files] + + if missing_files: + logger.warning(f"Stage 0 missed {len(missing_files)} files, adding to catch-all group") + catch_all_files = [ + ReviewFile(path=f, focus_areas=["general review"], risk_level="MEDIUM") + for f in missing_files + ] + plan.file_groups.append( + FileGroup( + group_id="uncategorized", + priority="MEDIUM", + rationale="Files not categorized by initial planning", + files=catch_all_files + ) + ) + + return plan diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/reconciliation.py b/python-ecosystem/mcp-client/service/review/orchestrator/reconciliation.py new file mode 100644 index 00000000..61d69b4d --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/reconciliation.py @@ -0,0 +1,284 @@ +""" +Issue reconciliation and deduplication logic for incremental reviews. +""" +import logging +from typing import Any, Dict, List, Optional + +from model.output_schemas import CodeReviewIssue + +logger = logging.getLogger(__name__) + + +def issue_matches_files(issue: Any, file_paths: List[str]) -> bool: + """Check if an issue is related to any of the given file paths. + + Matches on exact path, or when one path is a suffix of the other + (handles relative vs absolute paths). Does NOT match on basename alone + to avoid false positives (e.g., two different utils.py files). + """ + if hasattr(issue, 'model_dump'): + issue_data = issue.model_dump() + elif isinstance(issue, dict): + issue_data = issue + else: + issue_data = vars(issue) if hasattr(issue, '__dict__') else {} + + issue_file = issue_data.get('file', issue_data.get('filePath', '')) + if not issue_file: + return False + + for fp in file_paths: + if issue_file == fp or issue_file.endswith('/' + fp) or fp.endswith('/' + issue_file): + return True + return False + + +def compute_issue_fingerprint(data: dict) -> str: + """Compute a fingerprint for issue deduplication. + + Uses file + normalized line (±3 tolerance) + severity + truncated reason. + """ + file_path = data.get('file', data.get('filePath', '')) + line = data.get('line', data.get('lineNumber', 0)) + line_group = int(line) // 3 if line else 0 + severity = data.get('severity', '') + reason = data.get('reason', data.get('description', '')) + reason_prefix = reason[:50].lower().strip() if reason else '' + + return f"{file_path}::{line_group}::{severity}::{reason_prefix}" + + +def deduplicate_issues(issues: List[Any]) -> List[dict]: + """Deduplicate issues by fingerprint, keeping most recent version. + + If an older version is resolved but newer isn't, preserves resolved status. + """ + if not issues: + return [] + + deduped: dict = {} + + for issue in issues: + if hasattr(issue, 'model_dump'): + data = issue.model_dump() + elif isinstance(issue, dict): + data = issue.copy() + else: + data = vars(issue).copy() if hasattr(issue, '__dict__') else {} + + fingerprint = compute_issue_fingerprint(data) + existing = deduped.get(fingerprint) + + if existing is None: + deduped[fingerprint] = data + else: + existing_version = existing.get('prVersion') or 0 + current_version = data.get('prVersion') or 0 + existing_resolved = existing.get('status', '').lower() == 'resolved' + current_resolved = data.get('status', '').lower() == 'resolved' + + if current_version > existing_version: + # Current is newer + if existing_resolved and not current_resolved: + # Preserve resolved status from older version + data['status'] = 'resolved' + data['resolvedDescription'] = existing.get('resolvedDescription') + data['resolvedByCommit'] = existing.get('resolvedByCommit') + data['resolvedInPrVersion'] = existing.get('resolvedInPrVersion') + deduped[fingerprint] = data + elif current_version == existing_version: + # Same version - prefer resolved + if current_resolved and not existing_resolved: + deduped[fingerprint] = data + + return list(deduped.values()) + + +def format_previous_issues_for_batch(issues: List[Any]) -> str: + """Format previous issues for inclusion in batch prompt. + + Deduplicates issues first, then formats with resolution tracking so LLM knows: + - Which issues were previously found + - Which have been resolved (and how) + - Which PR version each issue was found/resolved in + """ + if not issues: + return "" + + # Deduplicate issues to avoid confusing the LLM with duplicates + deduped_issues = deduplicate_issues(issues) + + # Separate OPEN and RESOLVED issues + open_issues = [i for i in deduped_issues if i.get('status', '').lower() != 'resolved'] + resolved_issues = [i for i in deduped_issues if i.get('status', '').lower() == 'resolved'] + + lines = ["=== PREVIOUS ISSUES HISTORY (check if resolved/persisting) ==="] + lines.append("Issues have been deduplicated. Only check OPEN issues - RESOLVED ones are for context only.") + lines.append("") + + if open_issues: + lines.append("--- OPEN ISSUES (check if now fixed) ---") + for data in open_issues: + issue_id = data.get('id', 'unknown') + severity = data.get('severity', 'MEDIUM') + file_path = data.get('file', data.get('filePath', 'unknown')) + line = data.get('line', data.get('lineNumber', '?')) + reason = data.get('reason', data.get('description', 'No description')) + pr_version = data.get('prVersion', '?') + + lines.append(f"[ID:{issue_id}] {severity} @ {file_path}:{line} (v{pr_version})") + lines.append(f" Issue: {reason}") + lines.append("") + + if resolved_issues: + lines.append("--- RESOLVED ISSUES (for context only, do NOT re-report) ---") + for data in resolved_issues: + issue_id = data.get('id', 'unknown') + severity = data.get('severity', 'MEDIUM') + file_path = data.get('file', data.get('filePath', 'unknown')) + line = data.get('line', data.get('lineNumber', '?')) + reason = data.get('reason', data.get('description', 'No description')) + pr_version = data.get('prVersion', '?') + resolved_desc = data.get('resolvedDescription', '') + resolved_in = data.get('resolvedInPrVersion', '') + + lines.append(f"[ID:{issue_id}] {severity} @ {file_path}:{line} (v{pr_version}) - RESOLVED") + if resolved_desc: + lines.append(f" Resolution: {resolved_desc}") + if resolved_in: + lines.append(f" Resolved in: v{resolved_in}") + lines.append(f" Original issue: {reason}") + lines.append("") + + lines.append("INSTRUCTIONS:") + lines.append("- For OPEN issues that are now FIXED: report with 'isResolved': true (boolean)") + lines.append("- For OPEN issues still present: report with 'isResolved': false (boolean)") + lines.append("- Do NOT re-report RESOLVED issues - they are only shown for context") + lines.append("- IMPORTANT: 'isResolved' MUST be a JSON boolean (true/false), not a string") + lines.append("- Preserve the 'id' field for all issues you report from previous issues") + lines.append("=== END PREVIOUS ISSUES ===") + return "\n".join(lines) + + +async def reconcile_previous_issues( + request, + new_issues: List[CodeReviewIssue], + processed_diff = None +) -> List[CodeReviewIssue]: + """ + Reconcile previous issues with new findings in incremental mode. + - Mark resolved issues as isResolved=true + - Update line numbers for persisting issues + - Merge with new issues found in delta diff + - PRESERVE original issue data (reason, suggestedFixDescription, suggestedFixDiff) + """ + if not request.previousCodeAnalysisIssues: + return new_issues + + logger.info(f"Reconciling {len(request.previousCodeAnalysisIssues)} previous issues with {len(new_issues)} new issues") + + # Current commit for resolution tracking + current_commit = request.currentCommitHash or request.commitHash + + # Get the delta diff content to check what files/lines changed + delta_diff = request.deltaDiff or "" + + # Build a set of files that changed in the delta + changed_files_in_delta = set() + if processed_diff: + for f in processed_diff.files: + changed_files_in_delta.add(f.path) + + # Build lookup of previous issues by ID for merging with LLM results + prev_issues_by_id = {} + for prev_issue in request.previousCodeAnalysisIssues: + if hasattr(prev_issue, 'model_dump'): + prev_data = prev_issue.model_dump() + else: + prev_data = prev_issue if isinstance(prev_issue, dict) else vars(prev_issue) + issue_id = prev_data.get('id') + if issue_id: + prev_issues_by_id[str(issue_id)] = prev_data + + reconciled_issues = [] + processed_prev_ids = set() # Track which previous issues we've handled + + # Process new issues from LLM - merge with previous issue data if they reference same ID + for new_issue in new_issues: + new_data = new_issue.model_dump() if hasattr(new_issue, 'model_dump') else new_issue + issue_id = new_data.get('id') + + # If this issue references a previous issue ID, merge data + if issue_id and str(issue_id) in prev_issues_by_id: + prev_data = prev_issues_by_id[str(issue_id)] + processed_prev_ids.add(str(issue_id)) + + # Check if LLM marked it resolved + is_resolved = new_data.get('isResolved', False) + + # PRESERVE original data, use LLM's reason as resolution explanation + merged_issue = CodeReviewIssue( + id=str(issue_id), + severity=(prev_data.get('severity') or prev_data.get('issueSeverity') or 'MEDIUM').upper(), + category=prev_data.get('category') or prev_data.get('issueCategory') or prev_data.get('type') or 'CODE_QUALITY', + file=prev_data.get('file') or prev_data.get('filePath') or new_data.get('file', 'unknown'), + line=str(prev_data.get('line') or prev_data.get('lineNumber') or new_data.get('line', '1')), + # PRESERVE original reason and fix description + reason=prev_data.get('reason') or prev_data.get('title') or prev_data.get('description') or '', + suggestedFixDescription=prev_data.get('suggestedFixDescription') or prev_data.get('suggestedFix') or '', + suggestedFixDiff=prev_data.get('suggestedFixDiff') or None, + isResolved=is_resolved, + # Store LLM's explanation separately if resolved + resolutionExplanation=new_data.get('reason') if is_resolved else None, + resolvedInCommit=current_commit if is_resolved else None, + visibility=prev_data.get('visibility'), + codeSnippet=prev_data.get('codeSnippet') + ) + reconciled_issues.append(merged_issue) + else: + # New issue not referencing previous - keep as is + reconciled_issues.append(new_issue) + + # Process remaining previous issues not handled by LLM + for prev_issue in request.previousCodeAnalysisIssues: + if hasattr(prev_issue, 'model_dump'): + prev_data = prev_issue.model_dump() + else: + prev_data = prev_issue if isinstance(prev_issue, dict) else vars(prev_issue) + + issue_id = prev_data.get('id') + if issue_id and str(issue_id) in processed_prev_ids: + continue # Already handled above + + file_path = prev_data.get('file', prev_data.get('filePath', '')) + + # Check if this issue was already found in new issues (by file+line) + already_reported = False + for new_issue in new_issues: + new_data = new_issue.model_dump() if hasattr(new_issue, 'model_dump') else new_issue + if (new_data.get('file') == file_path and + str(new_data.get('line')) == str(prev_data.get('line', prev_data.get('lineNumber')))): + already_reported = True + break + + if already_reported: + continue + + # Preserve all original issue data + persisting_issue = CodeReviewIssue( + id=str(issue_id) if issue_id else None, + severity=(prev_data.get('severity') or prev_data.get('issueSeverity') or 'MEDIUM').upper(), + category=prev_data.get('category') or prev_data.get('issueCategory') or prev_data.get('type') or 'CODE_QUALITY', + file=file_path or prev_data.get('file') or prev_data.get('filePath') or 'unknown', + line=str(prev_data.get('line') or prev_data.get('lineNumber') or '1'), + reason=prev_data.get('reason') or prev_data.get('title') or prev_data.get('description') or '', + suggestedFixDescription=prev_data.get('suggestedFixDescription') or prev_data.get('suggestedFix') or '', + suggestedFixDiff=prev_data.get('suggestedFixDiff') or None, + isResolved=False, + visibility=prev_data.get('visibility'), + codeSnippet=prev_data.get('codeSnippet') + ) + reconciled_issues.append(persisting_issue) + + logger.info(f"Reconciliation complete: {len(reconciled_issues)} total issues") + return reconciled_issues diff --git a/python-ecosystem/mcp-client/service/review/orchestrator/stages.py b/python-ecosystem/mcp-client/service/review/orchestrator/stages.py new file mode 100644 index 00000000..90497c27 --- /dev/null +++ b/python-ecosystem/mcp-client/service/review/orchestrator/stages.py @@ -0,0 +1,926 @@ +""" +Stage execution methods for the multi-stage review pipeline. +""" +import json +import asyncio +import logging +import time +from typing import Any, Dict, List, Optional, Callable + +from model.dtos import ReviewRequestDto +from model.enrichment import PrEnrichmentDataDto +from model.output_schemas import CodeReviewOutput, CodeReviewIssue +from model.multi_stage import ( + ReviewPlan, + ReviewFile, + FileGroup, + CrossFileAnalysisResult, + FileReviewBatchOutput, +) +from utils.prompts.prompt_builder import PromptBuilder +from utils.diff_processor import ProcessedDiff, DiffProcessor +from utils.dependency_graph import create_smart_batches + +from service.review.orchestrator.agents import RecursiveMCPAgent, extract_llm_response_text +from service.review.orchestrator.json_utils import parse_llm_response +from service.review.orchestrator.reconciliation import ( + issue_matches_files, + format_previous_issues_for_batch, +) +from service.review.orchestrator.context_helpers import ( + extract_diff_snippets, + format_rag_context, +) + +logger = logging.getLogger(__name__) + + +async def execute_branch_analysis( + llm, + client, + prompt: str, + event_callback: Optional[Callable[[Dict], None]] = None +) -> Dict[str, Any]: + """ + Execute a single-pass branch analysis using the provided prompt. + """ + _emit_status(event_callback, "branch_analysis_started", "Starting Branch Analysis & Reconciliation...") + + agent = RecursiveMCPAgent( + llm=llm, + client=client, + additional_instructions=PromptBuilder.get_additional_instructions() + ) + + + try: + final_text = "" + # Branch analysis expects standard CodeReviewOutput + async for item in agent.stream(prompt, max_steps=15, output_schema=CodeReviewOutput): + if isinstance(item, CodeReviewOutput): + # Convert to dict format expected by service + issues = [i.model_dump() for i in item.issues] if item.issues else [] + return { + "issues": issues, + "comment": item.comment or "Branch analysis completed." + } + + if isinstance(item, str): + final_text = item + + # If stream finished without object, try parsing text + if final_text: + data = await parse_llm_response(final_text, CodeReviewOutput, llm) + issues = [i.model_dump() for i in data.issues] if data.issues else [] + return { + "issues": issues, + "comment": data.comment or "Branch analysis completed." + } + + return {"issues": [], "comment": "No issues found."} + + except Exception as e: + logger.error(f"Branch analysis failed: {e}", exc_info=True) + _emit_error(event_callback, str(e)) + raise + + +async def execute_stage_0_planning( + llm, + request: ReviewRequestDto, + is_incremental: bool = False, + processed_diff: Optional[ProcessedDiff] = None, +) -> ReviewPlan: + """ + Stage 0: Analyze metadata and generate a review plan. + Uses structured output for reliable JSON parsing. + """ + # Build a path → DiffFile lookup for real line stats + diff_by_path: Dict[str, Any] = {} + if processed_diff: + for df in processed_diff.files: + diff_by_path[df.path] = df + # Also index by basename for fuzzy matching + if '/' in df.path: + diff_by_path[df.path.rsplit('/', 1)[-1]] = df + + # Prepare context for prompt + changed_files_summary = [] + if request.changedFiles: + for f in request.changedFiles: + df = diff_by_path.get(f) or diff_by_path.get(f.rsplit('/', 1)[-1] if '/' in f else f) + changed_files_summary.append({ + "path": f, + "type": df.change_type.value.upper() if df else "MODIFIED", + "lines_added": df.additions if df else "?", + "lines_deleted": df.deletions if df else "?", + }) + + prompt = PromptBuilder.build_stage_0_planning_prompt( + repo_slug=request.projectVcsRepoSlug, + pr_id=str(request.pullRequestId), + pr_title=request.prTitle or "", + author="Unknown", + branch_name="source-branch", + target_branch=request.targetBranchName or "main", + commit_hash=request.commitHash or "HEAD", + changed_files_json=json.dumps(changed_files_summary, indent=2) + ) + + # Stage 0 uses direct LLM call (no tools needed - all metadata is provided) + try: + structured_llm = llm.with_structured_output(ReviewPlan) + result = await structured_llm.ainvoke(prompt) + if result: + logger.info("Stage 0 planning completed with structured output") + return result + except Exception as e: + logger.warning(f"Structured output failed for Stage 0: {e}") + + # Fallback to manual parsing + try: + response = await llm.ainvoke(prompt) + content = extract_llm_response_text(response) + return await parse_llm_response(content, ReviewPlan, llm) + except Exception as e: + logger.error(f"Stage 0 planning failed: {e}") + raise ValueError(f"Stage 0 planning failed: {e}") + + +def chunk_files(file_groups: List[Any], max_files_per_batch: int = 5) -> List[List[Dict[str, Any]]]: + """ + Flatten file groups and chunk into batches. + DEPRECATED: Use create_smart_batches for dependency-aware batching. + Kept for fallback when diff content is unavailable. + """ + all_files = [] + for group in file_groups: + for f in group.files: + # Attach priority context for the review + all_files.append({ + "file": f, + "priority": group.priority + }) + + return [all_files[i:i + max_files_per_batch] for i in range(0, len(all_files), max_files_per_batch)] + + +def create_smart_batches_wrapper( + file_groups: List[Any], + processed_diff: Optional[ProcessedDiff], + request: ReviewRequestDto, + rag_client, + max_files_per_batch: int = 7 +) -> List[List[Dict[str, Any]]]: + """ + Create dependency-aware batches that keep related files together. + + If enrichment data is available from Java (pre-computed relationships), + use it directly. Otherwise, use RAG's tree-sitter metadata to discover + file relationships: + - imports/exports: which files use symbols from other files + - class context: methods in the same class + - namespace context: files in the same package + + Falls back to directory-based grouping if both are unavailable. + """ + # Build branches list for RAG query + branches = [] + if request.targetBranchName: + branches.append(request.targetBranchName) + # Note: sourceBranchName is not in ReviewRequestDto, skip this check + if not branches: + branches = ['main', 'master'] # Fallback + + # Check for enrichment data from Java + enrichment_data = getattr(request, 'enrichmentData', None) + + try: + # Use RAG-based dependency analysis for intelligent batching + batches = create_smart_batches( + file_groups=file_groups, + workspace=request.projectWorkspace, + project=request.projectNamespace, + branches=branches, + rag_client=rag_client, + max_batch_size=max_files_per_batch, + enrichment_data=enrichment_data + ) + + # Log relationship analysis results + total_files = sum(len(b) for b in batches) + related_files = sum(1 for b in batches for f in b if f.get('has_relationships')) + enrichment_source = "enrichment data" if enrichment_data else "RAG discovery" + logger.info(f"Smart batching ({enrichment_source}): {total_files} files in {len(batches)} batches, " + f"{related_files} files have cross-file relationships") + + return batches + except Exception as e: + logger.warning(f"Smart batching failed, falling back to simple batching: {e}") + return chunk_files(file_groups, max_files_per_batch) + + +async def execute_stage_1_file_reviews( + llm, + request: ReviewRequestDto, + plan: ReviewPlan, + rag_client, + rag_context: Optional[Dict[str, Any]] = None, + processed_diff: Optional[ProcessedDiff] = None, + is_incremental: bool = False, + max_parallel: int = 5, + event_callback: Optional[Callable[[Dict], None]] = None, + pr_indexed: bool = False +) -> List[CodeReviewIssue]: + """ + Stage 1: Execute batch file reviews with per-batch RAG context. + Uses dependency-aware batching to keep related files together. + """ + # Use smart batching with RAG-based relationship discovery + batches = create_smart_batches_wrapper( + plan.file_groups, processed_diff, request, rag_client, max_files_per_batch=7 + ) + + total_files = sum(len(batch) for batch in batches) + related_batches = sum(1 for b in batches if any(f.get('has_relationships') for f in b)) + logger.info(f"Stage 1: Processing {total_files} files in {len(batches)} batches " + f"({related_batches} batches with cross-file relationships)") + + # Process batches with controlled parallelism + all_issues = [] + + # Process in waves to avoid rate limits + for wave_start in range(0, len(batches), max_parallel): + wave_end = min(wave_start + max_parallel, len(batches)) + wave_batches = batches[wave_start:wave_end] + wave_num = wave_start // max_parallel + 1 + + logger.info(f"Stage 1: Processing wave {wave_num}, " + f"batches {wave_start + 1}-{wave_end} of {len(batches)} IN PARALLEL") + + wave_start_time = time.time() + + tasks = [] + for batch_idx, batch in enumerate(wave_batches, start=wave_start + 1): + batch_paths = [item["file"].path for item in batch] + has_rels = any(item.get('has_relationships') for item in batch) + logger.debug(f"Batch {batch_idx}: {batch_paths} (cross-file relationships: {has_rels})") + # Create coroutine with batch_idx for tracking + tasks.append(_review_batch_with_timing( + batch_idx, llm, request, batch, rag_client, processed_diff, + is_incremental, rag_context, pr_indexed + )) + + # asyncio.gather runs all tasks CONCURRENTLY + results = await asyncio.gather(*tasks, return_exceptions=True) + + wave_elapsed = time.time() - wave_start_time + logger.info(f"Wave {wave_num} completed in {wave_elapsed:.2f}s " + f"({len(wave_batches)} batches parallel)") + + for idx, res in enumerate(results): + batch_num = wave_start + idx + 1 + if isinstance(res, Exception): + logger.error(f"Error reviewing batch {batch_num}: {res}") + elif res: + logger.info(f"Batch {batch_num} completed: {len(res)} issues found") + all_issues.extend(res) + else: + logger.info(f"Batch {batch_num} completed: no issues found") + + # Update progress + progress = 10 + int((wave_end / len(batches)) * 50) + _emit_progress(event_callback, progress, f"Stage 1: Reviewed {wave_end}/{len(batches)} batches") + + logger.info(f"Stage 1 Complete: {len(all_issues)} issues found across {total_files} files") + return all_issues + + +async def _review_batch_with_timing( + batch_idx: int, + llm, + request: ReviewRequestDto, + batch: List[Dict[str, Any]], + rag_client, + processed_diff: Optional[ProcessedDiff], + is_incremental: bool, + fallback_rag_context: Optional[Dict[str, Any]], + pr_indexed: bool +) -> List[CodeReviewIssue]: + """ + Wrapper that adds timing logs to show parallel execution. + """ + start_time = time.time() + batch_paths = [item["file"].path for item in batch] + logger.info(f"[Batch {batch_idx}] STARTED - files: {batch_paths}") + + try: + result = await review_file_batch( + llm, request, batch, rag_client, processed_diff, is_incremental, + fallback_rag_context=fallback_rag_context, pr_indexed=pr_indexed + ) + elapsed = time.time() - start_time + logger.info(f"[Batch {batch_idx}] FINISHED in {elapsed:.2f}s - {len(result)} issues") + return result + except Exception as e: + elapsed = time.time() - start_time + logger.error(f"[Batch {batch_idx}] FAILED after {elapsed:.2f}s: {e}") + raise + + +async def fetch_batch_rag_context( + rag_client, + request: ReviewRequestDto, + batch_file_paths: List[str], + batch_diff_snippets: List[str], + pr_indexed: bool = False +) -> Optional[Dict[str, Any]]: + """ + Fetch RAG context specifically for this batch of files. + + Two-pronged approach for comprehensive context: + 1. Semantic search using diff snippets (finds conceptually related code) + 2. Deterministic lookup using tree-sitter metadata (finds imported/referenced definitions) + + In hybrid mode (when PR files are indexed), passes pr_number to enable + queries that prioritize fresh PR data over potentially stale branch data. + """ + if not rag_client: + return None + + try: + # Determine branches for RAG query + rag_branch = request.targetBranchName or request.commitHash or "main" + base_branch = "main" # Default base branch for deterministic context + + logger.info(f"Fetching per-batch RAG context for {len(batch_file_paths)} files") + + # Use hybrid mode if PR files were indexed + pr_number = request.pullRequestId if pr_indexed else None + all_pr_files = request.changedFiles if pr_indexed else None + + # 1. Semantic search for conceptually related code + rag_response = await rag_client.get_pr_context( + workspace=request.projectWorkspace, + project=request.projectNamespace, + branch=rag_branch, + changed_files=batch_file_paths, + diff_snippets=batch_diff_snippets, + pr_title=request.prTitle, + pr_description=request.prDescription, + top_k=10, # Fewer chunks per batch for focused context + pr_number=pr_number, + all_pr_changed_files=all_pr_files + ) + + context = None + if rag_response and rag_response.get("context"): + context = rag_response.get("context") + chunk_count = len(context.get("relevant_code", [])) + logger.info(f"Semantic RAG: retrieved {chunk_count} chunks for batch") + + # 2. Deterministic lookup for cross-file dependencies (imports, extends, etc.) + # This uses tree-sitter metadata indexed during repo indexing + try: + deterministic_response = await rag_client.get_deterministic_context( + workspace=request.projectWorkspace, + project=request.projectNamespace, + branches=[rag_branch, base_branch], + file_paths=batch_file_paths, + limit_per_file=5 # Limit to avoid overwhelming context + ) + + if deterministic_response and deterministic_response.get("context"): + det_context = deterministic_response.get("context") + related_defs = det_context.get("related_definitions", {}) + + if related_defs: + # Merge deterministic results into the main context + if context is None: + context = {"relevant_code": []} + + # Add related definitions as additional context chunks + for def_name, def_chunks in related_defs.items(): + for chunk in def_chunks[:3]: # Limit per definition + context["relevant_code"].append({ + "file_path": chunk.get("file_path", ""), + "content": chunk.get("content", ""), + "score": 0.85, # High score for deterministic matches + "source": "deterministic", + "definition_name": def_name + }) + + logger.info(f"Deterministic RAG: added {len(related_defs)} related definitions") + + except Exception as det_err: + # Deterministic context is optional enhancement, don't fail the whole request + logger.debug(f"Deterministic RAG lookup skipped: {det_err}") + + if context: + total_chunks = len(context.get("relevant_code", [])) + logger.info(f"Total RAG context: {total_chunks} chunks for files {batch_file_paths}") + return context + + return None + + except Exception as e: + logger.warning(f"Failed to fetch per-batch RAG context: {e}") + return None + + +def _filter_rag_chunks_for_batch( + rag_context: Dict[str, Any], + batch_file_paths: List[str], +) -> Optional[Dict[str, Any]]: + """ + Pre-filter global RAG context to keep only chunks whose source path + is related to the current batch. This avoids injecting the full + 15-chunk global set into every batch when the per-batch fetch fails. + """ + chunks = rag_context.get("relevant_code", []) or rag_context.get("chunks", []) + if not chunks: + return rag_context + + batch_basenames = {p.rsplit("/", 1)[-1] if "/" in p else p for p in batch_file_paths} + batch_dirs = set() + for p in batch_file_paths: + parts = p.rsplit("/", 1) + if len(parts) == 2: + batch_dirs.add(parts[0]) + + filtered = [] + for chunk in chunks: + meta = chunk.get("metadata", {}) + chunk_path = meta.get("path") or chunk.get("path") or chunk.get("file_path", "") + if not chunk_path: + filtered.append(chunk) # keep chunks without path info + continue + + chunk_basename = chunk_path.rsplit("/", 1)[-1] if "/" in chunk_path else chunk_path + chunk_dir = chunk_path.rsplit("/", 1)[0] if "/" in chunk_path else "" + + # Keep if same file, same directory, or high-score (>= 0.8) regardless + score = chunk.get("score", chunk.get("relevance_score", 0)) + if (chunk_basename in batch_basenames + or chunk_dir in batch_dirs + or any(chunk_path.endswith(bp) or bp.endswith(chunk_path) for bp in batch_file_paths) + or score >= 0.8): + filtered.append(chunk) + + if not filtered: + # Don't return empty — keep original as fallback + return rag_context + + result = dict(rag_context) + result["relevant_code"] = filtered + return result + + +async def review_file_batch( + llm, + request: ReviewRequestDto, + batch_items: List[Dict[str, Any]], + rag_client, + processed_diff: Optional[ProcessedDiff] = None, + is_incremental: bool = False, + fallback_rag_context: Optional[Dict[str, Any]] = None, + pr_indexed: bool = False +) -> List[CodeReviewIssue]: + """ + Review a batch of files in a single LLM call with per-batch RAG context. + In incremental mode, uses delta diff and focuses on new changes only. + """ + batch_files_data = [] + batch_file_paths = [] + batch_diff_snippets = [] + #TODO: Project custom rules + project_rules = "" + + # For incremental mode, use deltaDiff instead of full diff + diff_source = None + if is_incremental and request.deltaDiff: + # Parse delta diff to extract per-file diffs + diff_source = DiffProcessor().process(request.deltaDiff) if request.deltaDiff else None + else: + diff_source = processed_diff + + # Collect file paths, diffs, and extract snippets for this batch + for item in batch_items: + file_info = item["file"] + batch_file_paths.append(file_info.path) + + # Extract diff from the appropriate source (delta for incremental, full for initial) + file_diff = "" + if diff_source: + for f in diff_source.files: + if f.path == file_info.path or f.path.endswith("/" + file_info.path): + file_diff = f.content + # Extract code snippets from diff for RAG semantic search + if file_diff: + batch_diff_snippets.extend(extract_diff_snippets(file_diff)) + break + + batch_files_data.append({ + "path": file_info.path, + "type": "MODIFIED", + "focus_areas": file_info.focus_areas, + "old_code": "", + "diff": file_diff or "(Diff unavailable)", + "is_incremental": is_incremental # Pass mode to prompt builder + }) + + # Fetch per-batch RAG context using batch-specific files and diff snippets + rag_context_text = "" + batch_rag_context = None + + if rag_client: + batch_rag_context = await fetch_batch_rag_context( + rag_client, request, batch_file_paths, batch_diff_snippets, pr_indexed + ) + + # Use batch-specific RAG context if available, otherwise fall back to initial context + # Hybrid mode: PR-indexed data is already included via fetch_batch_rag_context + if batch_rag_context: + logger.info(f"Using per-batch RAG context for: {batch_file_paths}") + rag_context_text = format_rag_context( + batch_rag_context, + set(batch_file_paths), + pr_changed_files=request.changedFiles + ) + elif fallback_rag_context: + # Filter the global RAG context to chunks relevant to this batch. + # Without filtering, every batch receives the same 15-chunk global set + # (wasting tokens on context unrelated to the batch's files). + filtered_fallback = _filter_rag_chunks_for_batch( + fallback_rag_context, batch_file_paths + ) + logger.info( + f"Using filtered fallback RAG context for batch: {batch_file_paths} " + f"({len((filtered_fallback or {}).get('relevant_code', []))} chunks)" + ) + rag_context_text = format_rag_context( + filtered_fallback, + set(batch_file_paths), + pr_changed_files=request.changedFiles + ) + + logger.info(f"RAG context for batch: {len(rag_context_text)} chars") + + # For incremental mode, filter previous issues relevant to this batch + # Also pass previous issues in FULL mode if they exist (subsequent PR iterations) + previous_issues_for_batch = "" + has_previous_issues = request.previousCodeAnalysisIssues and len(request.previousCodeAnalysisIssues) > 0 + if has_previous_issues: + relevant_prev_issues = [ + issue for issue in request.previousCodeAnalysisIssues + if issue_matches_files(issue, batch_file_paths) + ] + if relevant_prev_issues: + previous_issues_for_batch = format_previous_issues_for_batch(relevant_prev_issues) + + # Build ONE prompt for the batch with cross-file awareness + prompt = PromptBuilder.build_stage_1_batch_prompt( + files=batch_files_data, + priority=batch_items[0]["priority"] if batch_items else "MEDIUM", + project_rules=project_rules, + rag_context=rag_context_text, + is_incremental=is_incremental, + previous_issues=previous_issues_for_batch, + all_pr_files=request.changedFiles # Enable cross-file awareness in prompt + ) + + # Stage 1 uses direct LLM call (no tools needed - diff is already provided) + try: + # Try structured output first + structured_llm = llm.with_structured_output(FileReviewBatchOutput) + result = await structured_llm.ainvoke(prompt) + if result: + all_batch_issues = [] + for review in result.reviews: + all_batch_issues.extend(review.issues) + return all_batch_issues + except Exception as e: + logger.warning(f"Structured output failed for Stage 1 batch: {e}") + + # Fallback to manual parsing + try: + response = await llm.ainvoke(prompt) + content = extract_llm_response_text(response) + data = await parse_llm_response(content, FileReviewBatchOutput, llm) + all_batch_issues = [] + for review in data.reviews: + all_batch_issues.extend(review.issues) + return all_batch_issues + except Exception as parse_err: + logger.error( + f"Batch review double parse failure for {batch_file_paths}: {parse_err}. " + "Zero issues will be reported for this batch — results may be incomplete." + ) + return [] + + +# --------------------------------------------------------------------------- +# Stage 2 / Stage 3 context helpers +# --------------------------------------------------------------------------- + +# Simple substrings that mark a path as a migration / schema file. +# Checked with case-insensitive containment — no compiled regexes needed. +_MIGRATION_PATH_MARKERS = ( + '/db/migrate/', '/migrations/', '/migration/', + '/flyway/', '/liquibase/', '/alembic/', '/changeset/', +) + +# Fields to strip from CodeReviewIssue dicts before sending to Stage 2. +# Stage 2 only needs location + severity + reason to detect cross-file patterns. +_STAGE_2_STRIP_FIELDS = { + 'suggestedFixDiff', 'suggestedFixDescription', 'codeSnippet', + 'resolutionExplanation', 'resolvedInCommit', 'visibility', +} + + +def _build_architecture_context( + enrichment: Optional[PrEnrichmentDataDto], + changed_files: Optional[List[str]], +) -> str: + """ + Synthesise an architecture-reference section from the enrichment data + that pipeline-agent already computed (class hierarchy, inter-file + relationships, key imports). Zero extra LLM / RAG cost. + """ + sections: List[str] = [] + + if enrichment and enrichment.relationships: + rel_lines = [] + for r in enrichment.relationships: + rel_lines.append( + f" {r.sourceFile} --[{r.relationshipType.value}]--> {r.targetFile}" + + (f" (matched on: {r.matchedOn})" if r.matchedOn else "") + ) + if rel_lines: + sections.append( + "### Inter-file relationships (from dependency analysis)\n" + + "\n".join(rel_lines) + ) + + if enrichment and enrichment.fileMetadata: + hierarchy_lines = [] + for meta in enrichment.fileMetadata: + parts = [] + if meta.extendsClasses: + parts.append(f"extends {', '.join(meta.extendsClasses)}") + if meta.implementsInterfaces: + parts.append(f"implements {', '.join(meta.implementsInterfaces)}") + if parts: + hierarchy_lines.append(f" {meta.path}: {'; '.join(parts)}") + if hierarchy_lines: + sections.append( + "### Class hierarchy in changed files\n" + + "\n".join(hierarchy_lines) + ) + + # Summarise cross-file imports between changed files + if changed_files: + changed_set = set(changed_files or []) + import_lines = [] + for meta in enrichment.fileMetadata: + cross_imports = [ + imp for imp in meta.imports + if any(imp in cf or cf.endswith(imp) for cf in changed_set) + ] + if cross_imports: + import_lines.append( + f" {meta.path} imports: {', '.join(cross_imports[:10])}" + ) + if import_lines: + sections.append( + "### Cross-file imports among changed files\n" + + "\n".join(import_lines) + ) + + if not sections: + return "No architecture context available (enrichment data not provided)." + + return "\n\n".join(sections) + + +def _detect_migration_paths( + processed_diff: Optional[ProcessedDiff], +) -> str: + """ + Return a short list of migration file paths found in the diff. + + Stage 1 already reviews each migration file in full detail. + Stage 2 only needs to *know which files are migrations* so it can + reason about cross-file DB concerns (e.g. code referencing a column + that a migration drops). No raw SQL is injected. + """ + if not processed_diff: + return "No migration scripts detected." + + migration_files: List[str] = [] + for f in processed_diff.files: + path_lower = f.path.lower() + if path_lower.endswith('.sql') or any(m in path_lower for m in _MIGRATION_PATH_MARKERS): + migration_files.append(f.path) + + if not migration_files: + return "No migration scripts detected in this PR." + + listing = "\n".join(f"- {p}" for p in migration_files[:15]) # cap at 15 + return f"Migration files in this PR ({len(migration_files)}):\n{listing}" + + +def _slim_issues_for_stage_2( + issues: List[CodeReviewIssue], +) -> str: + """ + Serialize Stage 1 issues for Stage 2, stripping bulky fields that + Stage 2 does not need (fix diffs, code snippets, resolution details). + + Stage 2 detects *cross-file patterns* — it only needs: + file, line, severity, category, title/reason. + """ + slim = [] + for issue in issues: + d = issue.model_dump() + for key in _STAGE_2_STRIP_FIELDS: + d.pop(key, None) + slim.append(d) + return json.dumps(slim, indent=2) + + +def _summarize_issues_for_stage_3( + issues: List[CodeReviewIssue], +) -> str: + """ + Build a compact summary of Stage 1 issues for Stage 3 (executive report). + + The full issue list is posted as a separate comment, so Stage 3 only needs + aggregate counts and a short list of the most critical findings. + """ + if not issues: + return "No issues found in Stage 1." + + # --- Counts by severity --- + severity_counts: Dict[str, int] = {} + category_counts: Dict[str, int] = {} + for issue in issues: + sev = issue.severity.upper() + severity_counts[sev] = severity_counts.get(sev, 0) + 1 + cat = issue.category.upper() + category_counts[cat] = category_counts.get(cat, 0) + 1 + + lines = [ + f"Total issues: {len(issues)}", + "By severity: " + ", ".join(f"{k}: {v}" for k, v in sorted(severity_counts.items())), + "By category: " + ", ".join(f"{k}: {v}" for k, v in sorted(category_counts.items())), + ] + + # --- Top critical/high issues (title + file only) --- + priority_order = {'CRITICAL': 0, 'HIGH': 1, 'MEDIUM': 2, 'LOW': 3, 'INFO': 4} + ranked = sorted(issues, key=lambda i: priority_order.get(i.severity.upper(), 5)) + top_n = ranked[:10] + if top_n: + lines.append("\nTop findings:") + for i, issue in enumerate(top_n, 1): + lines.append(f" {i}. [{issue.severity}] {issue.file}: {issue.reason[:120]}") + + return "\n".join(lines) + + +async def execute_stage_2_cross_file( + llm, + request: ReviewRequestDto, + stage_1_issues: List[CodeReviewIssue], + plan: ReviewPlan, + processed_diff: Optional[ProcessedDiff] = None, +) -> CrossFileAnalysisResult: + """ + Stage 2: Cross-file analysis. + + Uses enrichment data (relationships, class hierarchy) and diff-detected + migrations to provide the LLM with real architecture context instead of + placeholders. + """ + # Slim Stage 1 findings (strip fix diffs, code snippets — Stage 2 only + # needs location + severity + reason for cross-file pattern detection) + issues_json = _slim_issues_for_stage_2(stage_1_issues) + + # Build architecture reference from enrichment data (zero-cost) + architecture_context = _build_architecture_context( + enrichment=request.enrichmentData, + changed_files=request.changedFiles, + ) + + # List migration file paths (no raw SQL — Stage 1 already reviewed them) + migrations = _detect_migration_paths(processed_diff) + + prompt = PromptBuilder.build_stage_2_cross_file_prompt( + repo_slug=request.projectVcsRepoSlug, + pr_title=request.prTitle or "", + commit_hash=request.commitHash or "HEAD", + stage_1_findings_json=issues_json, + architecture_context=architecture_context, + migrations=migrations, + cross_file_concerns=plan.cross_file_concerns + ) + + # Stage 2 uses direct LLM call (no tools needed - all data is provided from Stage 1) + try: + structured_llm = llm.with_structured_output(CrossFileAnalysisResult) + result = await structured_llm.ainvoke(prompt) + if result: + logger.info("Stage 2 cross-file analysis completed with structured output") + return result + except Exception as e: + logger.warning(f"Structured output failed for Stage 2: {e}") + + # Fallback to manual parsing + try: + response = await llm.ainvoke(prompt) + content = extract_llm_response_text(response) + return await parse_llm_response(content, CrossFileAnalysisResult, llm) + except Exception as e: + logger.error(f"Stage 2 cross-file analysis failed: {e}") + raise + + +async def execute_stage_3_aggregation( + llm, + request: ReviewRequestDto, + plan: ReviewPlan, + stage_1_issues: List[CodeReviewIssue], + stage_2_results: CrossFileAnalysisResult, + is_incremental: bool = False, + processed_diff: Optional[ProcessedDiff] = None +) -> str: + """ + Stage 3: Generate Markdown report. + In incremental mode, includes summary of resolved vs new issues. + """ + # Compact summary — the full issue list is posted as a separate comment + stage_1_json = _summarize_issues_for_stage_3(stage_1_issues) + stage_2_json = stage_2_results.model_dump_json(indent=2) + plan_json = plan.model_dump_json(indent=2) + + # Add incremental context to aggregation + incremental_context = "" + if is_incremental: + resolved_count = sum(1 for i in stage_1_issues if i.isResolved) + new_count = len(stage_1_issues) - resolved_count + previous_count = len(request.previousCodeAnalysisIssues or []) + incremental_context = f""" +## INCREMENTAL REVIEW SUMMARY +- Previous issues from last review: {previous_count} +- Issues resolved in this update: {resolved_count} +- New issues found in delta: {new_count} +- Total issues after reconciliation: {len(stage_1_issues)} +""" + + # Use real diff stats when available, fall back to 0 + additions = processed_diff.total_additions if processed_diff else 0 + deletions = processed_diff.total_deletions if processed_diff else 0 + + prompt = PromptBuilder.build_stage_3_aggregation_prompt( + repo_slug=request.projectVcsRepoSlug, + pr_id=str(request.pullRequestId), + author="Unknown", + pr_title=request.prTitle or "", + total_files=len(request.changedFiles or []), + additions=additions, + deletions=deletions, + stage_0_plan=plan_json, + stage_1_issues_json=stage_1_json, + stage_2_findings_json=stage_2_json, + recommendation=stage_2_results.pr_recommendation + ) + + response = await llm.ainvoke(prompt) + return extract_llm_response_text(response) + + +# Helper functions for event emission +def _emit_status(callback: Optional[Callable[[Dict], None]], state: str, message: str): + if callback: + callback({ + "type": "status", + "state": state, + "message": message + }) + + +def _emit_progress(callback: Optional[Callable[[Dict], None]], percent: int, message: str): + if callback: + callback({ + "type": "progress", + "percent": percent, + "message": message + }) + + +def _emit_error(callback: Optional[Callable[[Dict], None]], message: str): + if callback: + callback({ + "type": "error", + "message": message + }) diff --git a/python-ecosystem/mcp-client/service/review_service.py b/python-ecosystem/mcp-client/service/review/review_service.py similarity index 83% rename from python-ecosystem/mcp-client/service/review_service.py rename to python-ecosystem/mcp-client/service/review/review_service.py index ef9bb770..b3f4eb4e 100644 --- a/python-ecosystem/mcp-client/service/review_service.py +++ b/python-ecosystem/mcp-client/service/review/review_service.py @@ -6,18 +6,18 @@ from dotenv import load_dotenv from mcp_use import MCPClient -from model.models import ReviewRequestDto +from model.dtos import ReviewRequestDto from utils.mcp_config import MCPConfigBuilder from llm.llm_factory import LLMFactory from utils.prompts.prompt_builder import PromptBuilder from utils.response_parser import ResponseParser -from service.rag_client import RagClient, RAG_DEFAULT_TOP_K -from service.llm_reranker import LLMReranker -from service.issue_post_processor import post_process_analysis_result +from service.rag.rag_client import RagClient, RAG_DEFAULT_TOP_K +from service.rag.llm_reranker import LLMReranker +from service.review.issue_processor import post_process_analysis_result from utils.context_builder import (RAGMetrics, get_rag_cache) from utils.diff_processor import DiffProcessor from utils.error_sanitizer import create_user_friendly_error -from service.multi_stage_orchestrator import MultiStageReviewOrchestrator +from service.review.orchestrator import MultiStageReviewOrchestrator logger = logging.getLogger(__name__) @@ -30,8 +30,11 @@ class ReviewService: # Threshold for using LLM reranking (number of changed files) LLM_RERANK_FILE_THRESHOLD = 20 + # Maximum concurrent reviews (each spawns a JVM subprocess + LLM calls) + MAX_CONCURRENT_REVIEWS = int(os.environ.get("MAX_CONCURRENT_REVIEWS", "4")) + def __init__(self): - load_dotenv() + load_dotenv(interpolate=False) self.default_jar_path = os.environ.get( "MCP_SERVER_JAR", #"/var/www/html/codecrow/codecrow-public/java-ecosystem/mcp-servers/vcs-mcp/target/codecrow-vcs-mcp-1.0.jar", @@ -39,7 +42,7 @@ def __init__(self): ) self.rag_client = RagClient() self.rag_cache = get_rag_cache() - self.llm_reranker = None # Initialized lazily with LLM + self._review_semaphore = asyncio.Semaphore(self.MAX_CONCURRENT_REVIEWS) async def process_review_request( self, @@ -58,11 +61,12 @@ async def process_review_request( Returns: Dict with "result" key containing the analysis result or error """ - return await self._process_review( - request=request, - repo_path=None, - event_callback=event_callback - ) + async with self._review_semaphore: + return await self._process_review( + request=request, + repo_path=None, + event_callback=event_callback + ) async def _process_review( self, @@ -99,7 +103,8 @@ async def _process_review( has_raw_diff = bool(request.rawDiff) try: - context = "with pre-fetched diff" if has_raw_diff else "fetching diff via MCP" + async with asyncio.timeout(600): # 10-minute hard ceiling per review + context = "with pre-fetched diff" if has_raw_diff else "fetching diff via MCP" self._emit_event(event_callback, { "type": "status", "state": "started", @@ -120,9 +125,12 @@ async def _process_review( # Create LLM instance llm = self._create_llm(request) + + # Create a per-request reranker (not shared across concurrent requests) + llm_reranker = LLMReranker(llm_client=llm) # Fetch RAG context if enabled - rag_context = await self._fetch_rag_context(request, event_callback) + rag_context = await self._fetch_rag_context(request, event_callback, llm_reranker=llm_reranker) # Build processed_diff if rawDiff is available to optimize Stage 1 processed_diff = None @@ -164,22 +172,29 @@ async def _process_review( event_callback=event_callback ) - # Check for Branch Analysis / Reconciliation mode - if request.analysisType == "BRANCH_ANALYSIS": - logger.info("Executing Branch Analysis & Reconciliation mode") - # Build specific prompt for branch analysis - pr_metadata = self._build_pr_metadata(request) - prompt = PromptBuilder.build_branch_review_prompt_with_branch_issues_data(pr_metadata) - - result = await orchestrator.execute_branch_analysis(prompt) - else: - # Execute review with Multi-Stage Orchestrator - # Standard PR Review - result = await orchestrator.orchestrate_review( - request=request, - rag_context=rag_context, - processed_diff=processed_diff - ) + try: + # Check for Branch Analysis / Reconciliation mode + if request.analysisType == "BRANCH_ANALYSIS": + logger.info("Executing Branch Analysis & Reconciliation mode") + # Build specific prompt for branch analysis + pr_metadata = self._build_pr_metadata(request) + prompt = PromptBuilder.build_branch_review_prompt_with_branch_issues_data(pr_metadata) + + result = await orchestrator.execute_branch_analysis(prompt) + else: + # Execute review with Multi-Stage Orchestrator + # Standard PR Review + result = await orchestrator.orchestrate_review( + request=request, + rag_context=rag_context, + processed_diff=processed_diff + ) + finally: + # Always close MCP sessions to release JVM subprocesses + try: + await client.close_all_sessions() + except Exception as close_err: + logger.warning(f"Error closing MCP sessions: {close_err}") # Post-process issues to fix line numbers and merge duplicates @@ -221,6 +236,15 @@ async def _process_review( return {"result": result} + except TimeoutError: + timeout_msg = "Review timed out after 600 seconds" + logger.error(timeout_msg) + self._emit_event(event_callback, {"type": "error", "message": timeout_msg}) + error_response = ResponseParser.create_error_response( + "Review timed out", timeout_msg + ) + return {"result": error_response} + except Exception as e: # Log full error for debugging, but sanitize for user display logger.error(f"Review processing failed: {str(e)}", exc_info=True) @@ -256,7 +280,8 @@ def _build_jvm_props( async def _fetch_rag_context( self, request: ReviewRequestDto, - event_callback: Optional[Callable[[Dict], None]] + event_callback: Optional[Callable[[Dict], None]], + llm_reranker: Optional[LLMReranker] = None ) -> Optional[Dict[str, Any]]: """ Fetch relevant context from RAG pipeline. @@ -333,8 +358,8 @@ async def _fetch_rag_context( relevant_code = context.get("relevant_code", []) # Apply LLM reranking for large PRs - if len(changed_files) >= self.LLM_RERANK_FILE_THRESHOLD and self.llm_reranker: - reranked, rerank_result = await self.llm_reranker.rerank( + if len(changed_files) >= self.LLM_RERANK_FILE_THRESHOLD and llm_reranker: + reranked, rerank_result = await llm_reranker.rerank( relevant_code, pr_title=request.prTitle, pr_description=request.prDescription, @@ -391,17 +416,17 @@ def _create_mcp_client(self, config: Dict[str, Any]) -> MCPClient: raise Exception(f"Failed to construct MCPClient: {str(e)}") def _create_llm(self, request: ReviewRequestDto): - """Create LLM instance from request parameters and initialize reranker.""" + """Create LLM instance from request parameters.""" try: + # Log the model being used for this request + logger.info(f"Creating LLM for project {request.projectId}: provider={request.aiProvider}, model={request.aiModel}") + llm = LLMFactory.create_llm( request.aiModel, request.aiProvider, request.aiApiKey ) - # Initialize LLM reranker for large PRs - self.llm_reranker = LLMReranker(llm_client=llm) - return llm except Exception as e: raise Exception(f"Failed to create LLM instance: {str(e)}") diff --git a/python-ecosystem/mcp-client/tests/test_dependency_graph.py b/python-ecosystem/mcp-client/tests/test_dependency_graph.py new file mode 100644 index 00000000..9a06969e --- /dev/null +++ b/python-ecosystem/mcp-client/tests/test_dependency_graph.py @@ -0,0 +1,366 @@ +""" +Unit tests for the RAG-based dependency graph builder. +""" +import pytest +from unittest.mock import MagicMock, AsyncMock +from utils.dependency_graph import ( + DependencyGraphBuilder, + create_smart_batches, + FileNode, + FileRelationship +) + + +class MockReviewFile: + """Mock ReviewFile for testing.""" + def __init__(self, path: str, focus_areas: list = None): + self.path = path + self.focus_areas = focus_areas or [] + self.risk_level = "MEDIUM" + + +class MockFileGroup: + """Mock FileGroup for testing.""" + def __init__(self, files: list, priority: str = "MEDIUM"): + self.files = files + self.priority = priority + self.group_id = "test-group" + self.rationale = "Test group" + + +class MockRAGClient: + """Mock RAG client that returns predefined responses.""" + + def __init__(self, response: dict = None): + self.response = response or {} + self.call_count = 0 + self.last_call_args = None + + def get_deterministic_context(self, workspace, project, branches, file_paths, limit_per_file=15): + self.call_count += 1 + self.last_call_args = { + "workspace": workspace, + "project": project, + "branches": branches, + "file_paths": file_paths, + "limit_per_file": limit_per_file + } + return self.response + + +class TestDependencyGraphBuilder: + """Tests for DependencyGraphBuilder class.""" + + def test_init_without_rag_client(self): + """Test initialization without RAG client.""" + builder = DependencyGraphBuilder() + assert builder.rag_client is None + assert builder.nodes == {} + assert builder.relationships == [] + + def test_init_with_rag_client(self): + """Test initialization with RAG client.""" + mock_rag = MockRAGClient() + builder = DependencyGraphBuilder(rag_client=mock_rag) + assert builder.rag_client is mock_rag + + def test_build_basic_graph_fallback(self): + """Test fallback to basic graph when no RAG client.""" + builder = DependencyGraphBuilder() + + file_groups = [ + MockFileGroup([ + MockReviewFile("src/service/user_service.py"), + MockReviewFile("src/service/auth_service.py"), + ], priority="HIGH"), + MockFileGroup([ + MockReviewFile("tests/test_user.py"), + ], priority="MEDIUM"), + ] + + nodes = builder._build_basic_graph(file_groups) + + assert len(nodes) == 3 + assert "src/service/user_service.py" in nodes + # Files in same directory should be related + assert "src/service/auth_service.py" in nodes["src/service/user_service.py"].related_files + + def test_build_graph_from_rag_with_relationships(self): + """Test building graph with RAG-discovered relationships.""" + # Create RAG response with relationships + rag_response = { + "changed_files": { + "src/service/user_service.py": [ + { + "metadata": { + "primary_name": "UserService", + "semantic_names": ["get_user", "create_user"], + "imports": ["AuthService", "UserRepository"], + "parent_class": None, + "namespace": "src.service", + "path": "src/service/user_service.py" + } + } + ], + "src/service/auth_service.py": [ + { + "metadata": { + "primary_name": "AuthService", + "semantic_names": ["authenticate", "validate_token"], + "imports": [], + "parent_class": None, + "namespace": "src.service", + "path": "src/service/auth_service.py" + } + } + ] + }, + "related_definitions": { + "AuthService": [ + { + "metadata": { + "primary_name": "AuthService", + "path": "src/service/auth_service.py" + } + } + ] + }, + "class_context": {}, + "namespace_context": { + "src.service": [ + {"metadata": {"path": "src/service/user_service.py"}}, + {"metadata": {"path": "src/service/auth_service.py"}} + ] + } + } + + mock_rag = MockRAGClient(response=rag_response) + builder = DependencyGraphBuilder(rag_client=mock_rag) + + file_groups = [ + MockFileGroup([ + MockReviewFile("src/service/user_service.py"), + MockReviewFile("src/service/auth_service.py"), + ], priority="HIGH"), + ] + + nodes = builder.build_graph_from_rag( + file_groups, + workspace="test-workspace", + project="test-project", + branches=["main"] + ) + + assert len(nodes) == 2 + assert mock_rag.call_count == 1 + + # Check metadata extraction + user_node = nodes["src/service/user_service.py"] + assert "AuthService" in user_node.imports_symbols + assert "UserService" in user_node.exports_symbols + + # Check relationships discovered + assert len(builder.relationships) > 0 + + def test_get_connected_components(self): + """Test finding connected components.""" + builder = DependencyGraphBuilder() + + # Manually set up nodes with relationships + builder.nodes = { + "a.py": FileNode(path="a.py", priority="HIGH", related_files={"b.py"}), + "b.py": FileNode(path="b.py", priority="HIGH", related_files={"a.py", "c.py"}), + "c.py": FileNode(path="c.py", priority="HIGH", related_files={"b.py"}), + "d.py": FileNode(path="d.py", priority="MEDIUM", related_files=set()), # Isolated + } + + components = builder.get_connected_components() + + assert len(components) == 2 + # One component with a, b, c + abc_component = next(c for c in components if "a.py" in c) + assert abc_component == {"a.py", "b.py", "c.py"} + # One isolated component with d + d_component = next(c for c in components if "d.py" in c) + assert d_component == {"d.py"} + + def test_smart_batches_keeps_related_files_together(self): + """Test that smart batching keeps related files in same batch.""" + rag_response = { + "changed_files": { + "src/model/user.py": [ + {"metadata": {"primary_name": "User", "path": "src/model/user.py"}} + ], + "src/service/user_service.py": [ + {"metadata": { + "primary_name": "UserService", + "imports": ["User"], + "path": "src/service/user_service.py" + }} + ], + "src/api/user_api.py": [ + {"metadata": { + "primary_name": "UserAPI", + "imports": ["UserService"], + "path": "src/api/user_api.py" + }} + ], + "unrelated/config.py": [ + {"metadata": {"primary_name": "Config", "path": "unrelated/config.py"}} + ] + }, + "related_definitions": { + "User": [{"metadata": {"path": "src/model/user.py"}}], + "UserService": [{"metadata": {"path": "src/service/user_service.py"}}] + }, + "class_context": {}, + "namespace_context": {} + } + + mock_rag = MockRAGClient(response=rag_response) + builder = DependencyGraphBuilder(rag_client=mock_rag) + + file_groups = [ + MockFileGroup([ + MockReviewFile("src/model/user.py"), + MockReviewFile("src/service/user_service.py"), + MockReviewFile("src/api/user_api.py"), + MockReviewFile("unrelated/config.py"), + ], priority="HIGH"), + ] + + batches = builder.get_smart_batches( + file_groups, + workspace="test", + project="test", + branches=["main"], + max_batch_size=3 + ) + + # Should have batches created + assert len(batches) > 0 + + # Get all paths per batch + batch_paths = [[f['file'].path for f in b] for b in batches] + + # Related files should be in the same batch if possible + # (User -> UserService -> UserAPI chain) + for batch in batch_paths: + if "src/model/user.py" in batch: + # If User is in batch, UserService should be too (they're related) + if len(batch) > 1: + assert "src/service/user_service.py" in batch or "src/api/user_api.py" in batch + + def test_relationship_summary(self): + """Test getting relationship summary.""" + builder = DependencyGraphBuilder() + + builder.nodes = { + "a.py": FileNode(path="a.py", priority="HIGH", related_files={"b.py"}), + "b.py": FileNode(path="b.py", priority="HIGH", related_files={"a.py"}), + "c.py": FileNode(path="c.py", priority="LOW", related_files=set()), + } + builder.relationships = [ + FileRelationship("a.py", "b.py", "definition", "SomeClass", 0.95), + ] + + summary = builder.get_relationship_summary() + + assert summary["total_files"] == 3 + assert summary["total_relationships"] == 1 + assert summary["files_with_relationships"] == 2 + assert summary["relationship_types"]["definition"] == 1 + + +class TestCreateSmartBatches: + """Tests for the convenience function.""" + + def test_without_rag_client(self): + """Test smart batches without RAG client (fallback).""" + file_groups = [ + MockFileGroup([ + MockReviewFile("a.py"), + MockReviewFile("b.py"), + ], priority="HIGH"), + ] + + batches = create_smart_batches( + file_groups, + workspace="test", + project="test", + branches=["main"], + rag_client=None, + max_batch_size=5 + ) + + assert len(batches) > 0 + total_files = sum(len(b) for b in batches) + assert total_files == 2 + + def test_with_rag_client(self): + """Test smart batches with RAG client.""" + mock_rag = MockRAGClient(response={ + "changed_files": {}, + "related_definitions": {}, + "class_context": {}, + "namespace_context": {} + }) + + file_groups = [ + MockFileGroup([ + MockReviewFile("a.py"), + MockReviewFile("b.py"), + MockReviewFile("c.py"), + ], priority="HIGH"), + ] + + batches = create_smart_batches( + file_groups, + workspace="test", + project="test", + branches=["main"], + rag_client=mock_rag, + max_batch_size=2 + ) + + assert mock_rag.call_count == 1 + assert len(batches) >= 1 + + +class TestMergingSmallBatches: + """Tests for batch merging optimization.""" + + def test_merge_same_priority_batches(self): + """Test that small batches of same priority get merged.""" + builder = DependencyGraphBuilder() + + batches = [ + [{"file": MockReviewFile("a.py"), "priority": "HIGH"}], + [{"file": MockReviewFile("b.py"), "priority": "HIGH"}], + [{"file": MockReviewFile("c.py"), "priority": "LOW"}], + ] + + merged = builder._merge_small_batches(batches, min_size=2, max_size=5) + + # HIGH priority files should be merged + high_batch = next((b for b in merged if b[0]["priority"] == "HIGH"), None) + if high_batch: + assert len(high_batch) <= 5 + + def test_no_merge_when_at_max_size(self): + """Test that batches at max size are not merged.""" + builder = DependencyGraphBuilder() + + batches = [ + [ + {"file": MockReviewFile("a.py"), "priority": "HIGH"}, + {"file": MockReviewFile("b.py"), "priority": "HIGH"}, + {"file": MockReviewFile("c.py"), "priority": "HIGH"}, + ], + [{"file": MockReviewFile("d.py"), "priority": "HIGH"}], + ] + + merged = builder._merge_small_batches(batches, min_size=2, max_size=3) + + # First batch is already at max, so d.py stays separate + assert any(len(b) == 3 for b in merged) diff --git a/python-ecosystem/mcp-client/utils/context_builder.py b/python-ecosystem/mcp-client/utils/context_builder.py index e414d35e..9ebd1e1d 100644 --- a/python-ecosystem/mcp-client/utils/context_builder.py +++ b/python-ecosystem/mcp-client/utils/context_builder.py @@ -1,37 +1,18 @@ """ -Context builder for structured code review with Lost-in-the-Middle protection. -Implements priority-based context assembly with token budget management. +Context builder utilities for RAG metrics and caching. + +Note: The original ContextBuilder, ContextBudget, and MODEL_CONTEXT_LIMITS +were removed — they are superseded by the multi-stage orchestrator pipeline. """ import logging import os -import re import hashlib -from typing import Dict, List, Any, Optional, Tuple +from typing import Dict, List, Any, Optional from dataclasses import dataclass, field -from enum import Enum from datetime import datetime -from .file_classifier import FileClassifier, FilePriority, ClassifiedFile - logger = logging.getLogger(__name__) -# === Environment-based configuration === -# Context Budget percentages (must sum to 1.0) -CONTEXT_BUDGET_HIGH_PRIORITY_PCT = float(os.environ.get("CONTEXT_BUDGET_HIGH_PRIORITY_PCT", "0.30")) -CONTEXT_BUDGET_MEDIUM_PRIORITY_PCT = float(os.environ.get("CONTEXT_BUDGET_MEDIUM_PRIORITY_PCT", "0.40")) -CONTEXT_BUDGET_LOW_PRIORITY_PCT = float(os.environ.get("CONTEXT_BUDGET_LOW_PRIORITY_PCT", "0.20")) -CONTEXT_BUDGET_RAG_PCT = float(os.environ.get("CONTEXT_BUDGET_RAG_PCT", "0.10")) - -# Validate budget percentages sum to 1.0 -_budget_sum = CONTEXT_BUDGET_HIGH_PRIORITY_PCT + CONTEXT_BUDGET_MEDIUM_PRIORITY_PCT + CONTEXT_BUDGET_LOW_PRIORITY_PCT + CONTEXT_BUDGET_RAG_PCT -if abs(_budget_sum - 1.0) > 0.01: - logger.warning(f"Context budget percentages sum to {_budget_sum}, expected 1.0. Normalizing...") - _factor = 1.0 / _budget_sum - CONTEXT_BUDGET_HIGH_PRIORITY_PCT *= _factor - CONTEXT_BUDGET_MEDIUM_PRIORITY_PCT *= _factor - CONTEXT_BUDGET_LOW_PRIORITY_PCT *= _factor - CONTEXT_BUDGET_RAG_PCT *= _factor - # RAG Cache settings RAG_CACHE_TTL_SECONDS = int(os.environ.get("RAG_CACHE_TTL_SECONDS", "300")) RAG_CACHE_MAX_SIZE = int(os.environ.get("RAG_CACHE_MAX_SIZE", "100")) @@ -41,583 +22,6 @@ RAG_DEFAULT_TOP_K = int(os.environ.get("RAG_DEFAULT_TOP_K", "15")) -# Model context limits (approximate usable tokens after system prompt) -MODEL_CONTEXT_LIMITS = { - # OpenAI models - "gpt-4-turbo": 90000, - "gpt-4-turbo-preview": 90000, - "gpt-4-0125-preview": 90000, - "gpt-4-1106-preview": 90000, - "gpt-4": 6000, - "gpt-4-32k": 24000, - "gpt-3.5-turbo": 12000, - "gpt-3.5-turbo-16k": 12000, - # Anthropic models - "claude-3-opus": 140000, - "claude-3-opus-20240229": 140000, - "claude-3-sonnet": 140000, - "claude-3-sonnet-20240229": 140000, - "claude-3-haiku": 140000, - "claude-3-haiku-20240307": 140000, - "claude-3-5-sonnet": 140000, - "claude-3-5-sonnet-20240620": 140000, - "claude-3-5-sonnet-20241022": 140000, - # OpenRouter models (common ones) - "anthropic/claude-3-opus": 140000, - "anthropic/claude-3-sonnet": 140000, - "anthropic/claude-3-haiku": 140000, - "openai/gpt-4-turbo": 90000, - "openai/gpt-4": 6000, - "google/gemini-pro": 24000, - "google/gemini-pro-1.5": 700000, - "meta-llama/llama-3-70b-instruct": 6000, - "mistralai/mistral-large": 24000, - "deepseek/deepseek-chat": 48000, - "deepseek/deepseek-coder": 48000, - - "openai/gpt-5-mini": 400000, - "openai/gpt-5.1-codex-mini": 400000, - "openai/gpt-5.2": 512000, - "openai/gpt-5.1": 400000, - "openai/gpt-5.1-thinking": 196000, - "openai/gpt-5": 128000, - "openai/o3-high": 200000, - "openai/o4-mini": 128000, - - "anthropic/claude-4.5-opus": 500000, - "anthropic/claude-4.5-sonnet": 200000, - "anthropic/claude-4.5-haiku": 200000, - "anthropic/claude-3.7-sonnet": 200000, - - "google/gemini-3-pro": 1000000, - "google/gemini-3-flash": 1000000, - "google/gemini-2.5-pro": 2000000, - "google/gemini-2.5-flash": 1000000, - "google/gemini-3-flash-preview": 1000000, - "google/gemini-3-pro-preview": 1000000, - - "llama-4-405b": 128000, - "llama-4-70b": 128000, - "llama-4-scout": 1000000, - - "deepseek-v3.1": 160000, - "deepseek-v3": 128000, - "mistral-large-2025": 128000, - - "gpt-4o": 128000, - "gpt-4-turbo": 128000, - "claude-3-5-sonnet": 200000, - - "x-ai/grok-4.1-fast": 2000000, - - "default": 200000 -} - - -@dataclass -class ContextBudget: - """Token budget allocation for different context sections.""" - total_tokens: int = 45000 # Default max context budget - - # Budget distribution (percentages) - loaded from environment - high_priority_pct: float = field(default_factory=lambda: CONTEXT_BUDGET_HIGH_PRIORITY_PCT) - medium_priority_pct: float = field(default_factory=lambda: CONTEXT_BUDGET_MEDIUM_PRIORITY_PCT) - low_priority_pct: float = field(default_factory=lambda: CONTEXT_BUDGET_LOW_PRIORITY_PCT) - rag_context_pct: float = field(default_factory=lambda: CONTEXT_BUDGET_RAG_PCT) - - @property - def high_priority_tokens(self) -> int: - return int(self.total_tokens * self.high_priority_pct) - - @property - def medium_priority_tokens(self) -> int: - return int(self.total_tokens * self.medium_priority_pct) - - @property - def low_priority_tokens(self) -> int: - return int(self.total_tokens * self.low_priority_pct) - - @property - def rag_tokens(self) -> int: - return int(self.total_tokens * self.rag_context_pct) - - @classmethod - def for_model(cls, model_name: str) -> "ContextBudget": - """ - Create a ContextBudget optimized for a specific model. - - Args: - model_name: The model identifier (e.g., "gpt-4-turbo", "claude-3-opus") - - Returns: - ContextBudget with appropriate token limits - """ - # Normalize model name for lookup - model_lower = model_name.lower() - - # Try exact match first - limit = MODEL_CONTEXT_LIMITS.get(model_lower) - - # Try partial matches - if limit is None: - for key, value in MODEL_CONTEXT_LIMITS.items(): - if key in model_lower or model_lower in key: - limit = value - break - - # Use default if no match - if limit is None: - limit = MODEL_CONTEXT_LIMITS["default"] - logger.warning(f"Unknown model '{model_name}', using default context budget: {limit}") - - logger.info(f"Context budget for model '{model_name}': {limit} tokens") - return cls(total_tokens=limit) - - -@dataclass -class ContextSection: - """A section of structured context.""" - priority: str - title: str - description: str - content: str - file_count: int - token_estimate: int - files_included: List[str] = field(default_factory=list) - files_truncated: List[str] = field(default_factory=list) - - -@dataclass -class StructuredContext: - """Complete structured context with all sections.""" - sections: List[ContextSection] - total_files: int - total_tokens_estimated: int - files_analyzed: List[str] - files_skipped: List[str] - metadata: Dict[str, Any] = field(default_factory=dict) - - def to_prompt_string(self) -> str: - """Convert to formatted string for prompt injection.""" - parts = [] - - # Add metadata header - parts.append("=" * 60) - parts.append("STRUCTURED CONTEXT WITH PRIORITY MARKERS") - parts.append(f"Total files: {self.total_files} | Estimated tokens: {self.total_tokens_estimated}") - parts.append("=" * 60) - parts.append("") - - for section in self.sections: - if section.content: - parts.append(f"=== {section.priority} PRIORITY: {section.title} ===") - parts.append(f"({section.description})") - parts.append(f"Files: {section.file_count} | ~{section.token_estimate} tokens") - parts.append("") - parts.append(section.content) - parts.append("") - parts.append(f"=== END {section.priority} PRIORITY ===") - parts.append("") - - return "\n".join(parts) - - -class ContextBuilder: - """ - Builds structured context for code review with Lost-in-the-Middle protection. - - Key features: - 1. Priority-based file ordering (HIGH -> MEDIUM -> LOW -> RAG) - 2. Token budget management per section - 3. Explicit section markers for LLM attention - 4. Smart truncation when budget exceeded - """ - - # Approximate tokens per character (conservative estimate) - TOKENS_PER_CHAR = 0.25 - - def __init__(self, budget: Optional[ContextBudget] = None): - self.budget = budget or ContextBudget() - self.file_classifier = FileClassifier() - - def build_structured_context( - self, - pr_metadata: Dict[str, Any], - diff_content: Dict[str, str], # file_path -> diff/content - rag_context: Optional[Dict[str, Any]] = None, - file_paths: Optional[List[str]] = None - ) -> StructuredContext: - """ - Build structured context from PR data and RAG results. - - Args: - pr_metadata: PR metadata (title, description, etc.) - diff_content: Dict mapping file paths to their diff content - rag_context: Optional RAG query results - file_paths: Optional list of changed file paths - - Returns: - StructuredContext with prioritized sections - """ - file_paths = file_paths or list(diff_content.keys()) - - # Step 1: Classify files by priority - classified = FileClassifier.classify_files(file_paths) - stats = FileClassifier.get_priority_stats(classified) - logger.info(f"File classification: {stats}") - - sections = [] - files_analyzed = [] - files_skipped = [] - total_tokens = 0 - - # Step 2: Build HIGH priority section - high_section = self._build_priority_section( - priority=FilePriority.HIGH, - files=classified[FilePriority.HIGH], - diff_content=diff_content, - token_budget=self.budget.high_priority_tokens, - title="Core Business Logic", - description="Analyze FIRST - security, auth, core services" - ) - sections.append(high_section) - files_analyzed.extend(high_section.files_included) - files_skipped.extend(high_section.files_truncated) - total_tokens += high_section.token_estimate - - # Step 3: Build MEDIUM priority section - medium_section = self._build_priority_section( - priority=FilePriority.MEDIUM, - files=classified[FilePriority.MEDIUM], - diff_content=diff_content, - token_budget=self.budget.medium_priority_tokens, - title="Dependencies & Shared Utils", - description="Models, DTOs, utilities, components" - ) - sections.append(medium_section) - files_analyzed.extend(medium_section.files_included) - files_skipped.extend(medium_section.files_truncated) - total_tokens += medium_section.token_estimate - - # Step 4: Build LOW priority section - low_section = self._build_priority_section( - priority=FilePriority.LOW, - files=classified[FilePriority.LOW], - diff_content=diff_content, - token_budget=self.budget.low_priority_tokens, - title="Tests & Config", - description="Test files, configurations (quick scan)" - ) - sections.append(low_section) - files_analyzed.extend(low_section.files_included) - files_skipped.extend(low_section.files_truncated) - total_tokens += low_section.token_estimate - - # Step 5: Build RAG context section - rag_section = self._build_rag_section( - rag_context=rag_context, - token_budget=self.budget.rag_tokens, - already_included=set(files_analyzed) - ) - sections.append(rag_section) - total_tokens += rag_section.token_estimate - - # Add skipped files from classification - for f in classified[FilePriority.SKIP]: - files_skipped.append(f.path) - - return StructuredContext( - sections=sections, - total_files=len(file_paths), - total_tokens_estimated=total_tokens, - files_analyzed=files_analyzed, - files_skipped=files_skipped, - metadata={ - "classification_stats": stats, - "budget": { - "total": self.budget.total_tokens, - "high": self.budget.high_priority_tokens, - "medium": self.budget.medium_priority_tokens, - "low": self.budget.low_priority_tokens, - "rag": self.budget.rag_tokens - } - } - ) - - def _build_priority_section( - self, - priority: FilePriority, - files: List[ClassifiedFile], - diff_content: Dict[str, str], - token_budget: int, - title: str, - description: str - ) -> ContextSection: - """Build a single priority section with budget management.""" - content_parts = [] - files_included = [] - files_truncated = [] - current_tokens = 0 - - for classified_file in files: - path = classified_file.path - if path not in diff_content: - continue - - file_content = diff_content[path] - file_tokens = self._estimate_tokens(file_content) - - # Check if adding this file would exceed budget - if current_tokens + file_tokens > token_budget: - # Try to include truncated version if significant budget remains - remaining_budget = token_budget - current_tokens - if remaining_budget > 500: # Only truncate if meaningful space left - truncated_content = self._truncate_content(file_content, remaining_budget) - content_parts.append(f"### {path} (TRUNCATED)") - content_parts.append(truncated_content) - content_parts.append("") - files_included.append(path) - current_tokens += self._estimate_tokens(truncated_content) - else: - files_truncated.append(path) - continue - - content_parts.append(f"### {path}") - content_parts.append(f"Category: {classified_file.category} | Importance: {classified_file.estimated_importance:.2f}") - content_parts.append("```") - content_parts.append(file_content) - content_parts.append("```") - content_parts.append("") - - files_included.append(path) - current_tokens += file_tokens - - return ContextSection( - priority=priority.value, - title=title, - description=description, - content="\n".join(content_parts), - file_count=len(files_included), - token_estimate=current_tokens, - files_included=files_included, - files_truncated=files_truncated - ) - - def _build_rag_section( - self, - rag_context: Optional[Dict[str, Any]], - token_budget: int, - already_included: set - ) -> ContextSection: - """Build RAG context section with deduplication.""" - if not rag_context or not rag_context.get("relevant_code"): - return ContextSection( - priority="RAG", - title="Additional Context from Codebase", - description="No RAG context available", - content="", - file_count=0, - token_estimate=0 - ) - - content_parts = [] - files_included = [] - current_tokens = 0 - max_rag_chunks = 5 # Limit RAG chunks - - relevant_code = rag_context.get("relevant_code", []) - chunk_count = 0 - - for chunk in relevant_code: - if chunk_count >= max_rag_chunks: - break - - # Skip if file already included in priority sections - chunk_path = chunk.get("metadata", {}).get("path", "unknown") - if chunk_path in already_included: - continue - - chunk_text = chunk.get("text", "") - chunk_score = chunk.get("score", 0) - chunk_tokens = self._estimate_tokens(chunk_text) - - if current_tokens + chunk_tokens > token_budget: - break - - content_parts.append(f"### RAG Context {chunk_count + 1}: {chunk_path}") - content_parts.append(f"Relevance score: {chunk_score:.3f}") - content_parts.append("```") - content_parts.append(chunk_text) - content_parts.append("```") - content_parts.append("") - - files_included.append(chunk_path) - current_tokens += chunk_tokens - chunk_count += 1 - - return ContextSection( - priority="RAG", - title="Additional Context from Codebase", - description="Semantically relevant code from repository (max 5 chunks)", - content="\n".join(content_parts), - file_count=len(files_included), - token_estimate=current_tokens, - files_included=files_included - ) - - def _estimate_tokens(self, text: str) -> int: - """Estimate token count for text.""" - return int(len(text) * self.TOKENS_PER_CHAR) - - def _truncate_content(self, content: str, target_tokens: int) -> str: - """Truncate content to fit within token budget.""" - target_chars = int(target_tokens / self.TOKENS_PER_CHAR) - - if len(content) <= target_chars: - return content - - # Try to truncate at a logical boundary - truncated = content[:target_chars] - - # Find last newline to avoid cutting mid-line - last_newline = truncated.rfind('\n') - if last_newline > target_chars * 0.8: # Only use if we don't lose too much - truncated = truncated[:last_newline] - - return truncated + "\n... [TRUNCATED - remaining content omitted for token budget]" - - -class RagReranker: - """ - Reranks RAG results for better relevance using multiple strategies. - """ - - @staticmethod - def rerank_by_file_priority( - rag_results: List[Dict[str, Any]], - classified_files: Dict[FilePriority, List[ClassifiedFile]], - boost_factor: float = 1.5 - ) -> List[Dict[str, Any]]: - """ - Boost RAG results that relate to high-priority files. - - Args: - rag_results: Original RAG results - classified_files: Files classified by priority - boost_factor: Multiplier for scores of related high-priority results - - Returns: - Reranked results with adjusted scores - """ - # Build set of high-priority file paths - high_priority_paths = { - f.path for f in classified_files.get(FilePriority.HIGH, []) - } - - # Get directory patterns from high-priority files - high_priority_dirs = set() - for path in high_priority_paths: - parts = path.split('/') - for i in range(1, len(parts)): - high_priority_dirs.add('/'.join(parts[:i])) - - reranked = [] - for result in rag_results: - result_copy = result.copy() - result_path = result.get("metadata", {}).get("path", "") - - # Boost if directly matches high-priority file - if result_path in high_priority_paths: - result_copy["score"] = result.get("score", 0) * boost_factor - result_copy["_boost_reason"] = "direct_high_priority_match" - # Smaller boost if in high-priority directory - elif any(result_path.startswith(d) for d in high_priority_dirs): - result_copy["score"] = result.get("score", 0) * (boost_factor * 0.7) - result_copy["_boost_reason"] = "high_priority_directory" - - reranked.append(result_copy) - - # Sort by adjusted score - reranked.sort(key=lambda x: x.get("score", 0), reverse=True) - return reranked - - @staticmethod - def filter_by_relevance_threshold( - results: List[Dict[str, Any]], - min_score: float = None, - min_results: int = 3 - ) -> List[Dict[str, Any]]: - """ - Filter results by minimum relevance score. - Always returns at least min_results if available. - Default min_score from RAG_MIN_RELEVANCE_SCORE env var. - """ - if min_score is None: - min_score = RAG_MIN_RELEVANCE_SCORE - filtered = [r for r in results if r.get("score", 0) >= min_score] - - # Ensure minimum results - if len(filtered) < min_results and len(results) >= min_results: - # Add top results regardless of score - for result in results: - if result not in filtered: - filtered.append(result) - if len(filtered) >= min_results: - break - - return filtered - - @staticmethod - def deduplicate_by_content( - results: List[Dict[str, Any]], - similarity_threshold: float = 0.8 - ) -> List[Dict[str, Any]]: - """ - Remove near-duplicate results based on content similarity. - Uses simple text overlap for efficiency. - """ - if not results: - return [] - - unique_results = [results[0]] - - for result in results[1:]: - is_duplicate = False - result_text = result.get("text", "") - - for existing in unique_results: - existing_text = existing.get("text", "") - - # Simple overlap check - overlap = RagReranker._calculate_overlap(result_text, existing_text) - if overlap > similarity_threshold: - is_duplicate = True - break - - if not is_duplicate: - unique_results.append(result) - - return unique_results - - @staticmethod - def _calculate_overlap(text1: str, text2: str) -> float: - """Calculate simple text overlap ratio.""" - if not text1 or not text2: - return 0.0 - - # Use set of words for simple comparison - words1 = set(text1.lower().split()) - words2 = set(text2.lower().split()) - - if not words1 or not words2: - return 0.0 - - intersection = len(words1 & words2) - union = len(words1 | words2) - - return intersection / union if union > 0 else 0.0 - - @dataclass class RAGMetrics: """Metrics for RAG processing quality and performance.""" diff --git a/python-ecosystem/mcp-client/utils/dependency_graph.py b/python-ecosystem/mcp-client/utils/dependency_graph.py new file mode 100644 index 00000000..e8d3cd0e --- /dev/null +++ b/python-ecosystem/mcp-client/utils/dependency_graph.py @@ -0,0 +1,630 @@ +""" +Dependency graph builder for intelligent file batching. + +SMART APPROACH: Leverages RAG's pre-indexed tree-sitter metadata to discover +file relationships instead of re-parsing diffs with regex. + +The RAG system already has: +- semantic_names: function/method/class names +- imports: import statements parsed by tree-sitter +- extends: parent classes/interfaces +- parent_class: containing class +- namespace: package/namespace + +This module queries RAG to build a relationship graph, enabling intelligent +batching that keeps related files together for better cross-file context. +""" +import logging +from collections import defaultdict +from typing import Dict, List, Set, Any, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from service.rag.rag_client import RagClient + +from dataclasses import dataclass, field + +logger = logging.getLogger(__name__) + + +@dataclass +class FileNode: + """Represents a file in the dependency graph.""" + path: str + priority: str + # Relationships discovered from RAG tree-sitter metadata + related_files: Set[str] = field(default_factory=set) + imports_symbols: Set[str] = field(default_factory=set) + exports_symbols: Set[str] = field(default_factory=set) + parent_classes: Set[str] = field(default_factory=set) + namespaces: Set[str] = field(default_factory=set) + extends: Set[str] = field(default_factory=set) + focus_areas: List[str] = field(default_factory=list) + relationship_strength: float = 0.0 + + +@dataclass +class FileRelationship: + """Represents a relationship between two files.""" + source_file: str + target_file: str + relationship_type: str # 'definition', 'same_class', 'same_namespace' + matched_on: str + strength: float + + +class DependencyGraphBuilder: + """ + Builds a dependency graph using RAG's tree-sitter metadata or pre-computed relationships. + + SMART APPROACH (v2): When Java sends enrichment data with pre-computed relationships, + use those directly instead of querying RAG. This eliminates duplicate work since + Java already called RAG's /parse endpoint. + + Fallback: When no enrichment data available, query RAG's deterministic context API + which has the FULL file indexed with tree-sitter metadata. + + Relationship types discovered: + - definition: File A uses symbol defined in File B + - same_class: Files contain methods of the same class + - same_namespace: Files are in the same package/namespace + """ + + RELATIONSHIP_WEIGHTS = { + 'changed_file': 1.0, + 'definition': 0.95, + 'IMPORTS': 0.90, + 'EXTENDS': 0.95, + 'IMPLEMENTS': 0.95, + 'CALLS': 0.85, + 'class_context': 0.85, + 'namespace_context': 0.75, + 'SAME_PACKAGE': 0.60, + } + + def __init__(self, rag_client: Optional["RAGClient"] = None): + self.rag_client = rag_client + self.nodes: Dict[str, FileNode] = {} + self.relationships: List[FileRelationship] = [] + self._metadata_cache: Dict[str, Dict] = {} + + def build_graph_from_enrichment( + self, + file_groups: List[Any], + enrichment_data: Any, + ) -> Dict[str, FileNode]: + """ + Build dependency graph from pre-computed enrichment data sent by Java. + + This is the preferred method when enrichment_data is available, as it: + - Eliminates redundant RAG calls (Java already parsed the files) + - Uses full file content for accurate relationship detection + - Provides relationships computed with proper AST parsing + + Args: + file_groups: List of FileGroup objects with files to analyze + enrichment_data: PrEnrichmentDataDto from Java with relationships and metadata + + Returns: + Dict of file paths to FileNode objects with relationships populated + """ + if not enrichment_data or not enrichment_data.has_data(): + logger.info("No enrichment data available, falling back to basic grouping") + return self._build_basic_graph(file_groups) + + # Initialize nodes from file groups + for group in file_groups: + for f in group.files: + self.nodes[f.path] = FileNode( + path=f.path, + priority=group.priority, + focus_areas=f.focus_areas if hasattr(f, 'focus_areas') else [] + ) + + # Process pre-computed relationships + relationships_by_file: Dict[str, Set[str]] = defaultdict(set) + + for rel in enrichment_data.relationships: + source = rel.sourceFile + target = rel.targetFile + rel_type = rel.relationshipType.value if hasattr(rel.relationshipType, 'value') else str(rel.relationshipType) + + # Only add relationships between files we're analyzing + if source in self.nodes and target in self.nodes: + relationships_by_file[source].add(target) + relationships_by_file[target].add(source) + + weight = self.RELATIONSHIP_WEIGHTS.get(rel_type, 0.5) + self.relationships.append(FileRelationship( + source_file=source, + target_file=target, + relationship_type=rel_type, + matched_on=rel.matchedOn or "", + strength=weight + )) + + # Process metadata to populate node symbols + for meta in enrichment_data.fileMetadata: + if meta.path in self.nodes: + node = self.nodes[meta.path] + if meta.imports: + node.imports_symbols.update(meta.imports) + if meta.semanticNames: + node.exports_symbols.update(meta.semanticNames) + if meta.extendsClasses: + node.extends.update(meta.extendsClasses) + if meta.parentClass: + node.parent_classes.add(meta.parentClass) + if meta.namespace: + node.namespaces.add(meta.namespace) + + # Update nodes with discovered relationships + for file_path, related in relationships_by_file.items(): + if file_path in self.nodes: + self.nodes[file_path].related_files.update(related) + self.nodes[file_path].relationship_strength = self._calculate_strength( + file_path, related + ) + + logger.info( + f"Dependency graph built from enrichment: {len(self.nodes)} files, " + f"{len(self.relationships)} relationships" + ) + + return self.nodes + + def build_graph_from_rag( + self, + file_groups: List[Any], + workspace: str, + project: str, + branches: List[str], + ) -> Dict[str, FileNode]: + """ + Build dependency graph by querying RAG's deterministic context API. + + This leverages tree-sitter metadata extracted during indexing: + - imports, extends, parent_class, namespace, semantic_names + """ + if not self.rag_client: + logger.warning("No RAG client provided, falling back to basic grouping") + return self._build_basic_graph(file_groups) + + # Collect all file paths + all_file_paths = [] + file_priority_map = {} + file_info_map = {} + + for group in file_groups: + for f in group.files: + all_file_paths.append(f.path) + file_priority_map[f.path] = group.priority + file_info_map[f.path] = f + self.nodes[f.path] = FileNode( + path=f.path, + priority=group.priority, + focus_areas=f.focus_areas if hasattr(f, 'focus_areas') else [] + ) + + if not all_file_paths: + return self.nodes + + # Query RAG for deterministic context + try: + rag_response = self.rag_client.get_deterministic_context( + workspace=workspace, + project=project, + branches=branches, + file_paths=all_file_paths, + limit_per_file=15 + ) + self._metadata_cache['last_response'] = rag_response + except Exception as e: + logger.warning(f"RAG query failed, falling back to basic grouping: {e}") + return self._build_basic_graph(file_groups) + + # Extract relationships from RAG response + self._extract_relationships_from_rag(rag_response, all_file_paths) + + logger.info( + f"Dependency graph built: {len(self.nodes)} files, " + f"{len(self.relationships)} relationships" + ) + + return self.nodes + + def _extract_relationships_from_rag( + self, + rag_response: Dict, + changed_file_paths: List[str] + ) -> None: + """Extract file relationships from RAG deterministic context response.""" + changed_file_set = set(changed_file_paths) + file_relationships: Dict[str, Set[str]] = defaultdict(set) + + # Process changed_files to extract metadata + changed_files = rag_response.get('changed_files', {}) + for file_path, chunks in changed_files.items(): + norm_path = file_path.lstrip('/') + if norm_path in self.nodes: + for chunk in chunks: + metadata = chunk.get('metadata', {}) + + # Extract symbols this file exports (defines) + if metadata.get('primary_name'): + self.nodes[norm_path].exports_symbols.add(metadata['primary_name']) + if metadata.get('semantic_names'): + self.nodes[norm_path].exports_symbols.update(metadata['semantic_names']) + + # Extract what this file imports + if metadata.get('imports'): + for imp in metadata['imports']: + if isinstance(imp, str): + parts = imp.replace(';', '').split('\\') + if parts: + self.nodes[norm_path].imports_symbols.add(parts[-1].strip()) + + # Track class/namespace membership + if metadata.get('parent_class'): + self.nodes[norm_path].parent_classes.add(metadata['parent_class']) + if metadata.get('namespace'): + self.nodes[norm_path].namespaces.add(metadata['namespace']) + if metadata.get('extends'): + self.nodes[norm_path].extends.update(metadata['extends']) + + # Process related_definitions + related_definitions = rag_response.get('related_definitions', {}) + for symbol, chunks in related_definitions.items(): + for chunk in chunks: + metadata = chunk.get('metadata', {}) + related_path = metadata.get('path', '').lstrip('/') + + if related_path and related_path in self.nodes: + for file_path in changed_file_set: + norm_path = file_path.lstrip('/') + if norm_path in self.nodes: + node = self.nodes[norm_path] + if symbol in node.imports_symbols or symbol in node.exports_symbols: + file_relationships[norm_path].add(related_path) + file_relationships[related_path].add(norm_path) + self.relationships.append(FileRelationship( + source_file=norm_path, + target_file=related_path, + relationship_type='definition', + matched_on=symbol, + strength=self.RELATIONSHIP_WEIGHTS['definition'] + )) + + # Process class_context + class_context = rag_response.get('class_context', {}) + for parent_class, chunks in class_context.items(): + class_files = set() + for chunk in chunks: + metadata = chunk.get('metadata', {}) + related_path = metadata.get('path', '').lstrip('/') + if related_path in self.nodes: + class_files.add(related_path) + + for f1 in class_files: + for f2 in class_files: + if f1 != f2: + file_relationships[f1].add(f2) + if f1 < f2: + self.relationships.append(FileRelationship( + source_file=f1, + target_file=f2, + relationship_type='same_class', + matched_on=parent_class, + strength=self.RELATIONSHIP_WEIGHTS['class_context'] + )) + + # Process namespace_context + namespace_context = rag_response.get('namespace_context', {}) + for namespace, chunks in namespace_context.items(): + ns_files = set() + for chunk in chunks: + metadata = chunk.get('metadata', {}) + related_path = metadata.get('path', '').lstrip('/') + if related_path in self.nodes: + ns_files.add(related_path) + + for f1 in ns_files: + for f2 in ns_files: + if f1 != f2: + file_relationships[f1].add(f2) + if f1 < f2: + self.relationships.append(FileRelationship( + source_file=f1, + target_file=f2, + relationship_type='same_namespace', + matched_on=namespace, + strength=self.RELATIONSHIP_WEIGHTS['namespace_context'] + )) + + # Update nodes with discovered relationships + for file_path, related in file_relationships.items(): + if file_path in self.nodes: + self.nodes[file_path].related_files.update(related) + self.nodes[file_path].relationship_strength = self._calculate_strength( + file_path, related + ) + + def _calculate_strength(self, file_path: str, related_files: Set[str]) -> float: + total_strength = 0.0 + for rel in self.relationships: + if rel.source_file == file_path or rel.target_file == file_path: + total_strength += rel.strength + return min(total_strength, 5.0) + + def _build_basic_graph(self, file_groups: List[Any]) -> Dict[str, FileNode]: + """Fallback: build basic graph without RAG (by directory).""" + for group in file_groups: + for f in group.files: + self.nodes[f.path] = FileNode( + path=f.path, + priority=group.priority, + focus_areas=f.focus_areas if hasattr(f, 'focus_areas') else [] + ) + + # Files in same directory are related + dir_files: Dict[str, List[str]] = defaultdict(list) + for path in self.nodes: + dir_path = '/'.join(path.split('/')[:-1]) if '/' in path else '' + dir_files[dir_path].append(path) + + for dir_path, files in dir_files.items(): + if len(files) > 1: + for f1 in files: + for f2 in files: + if f1 != f2: + self.nodes[f1].related_files.add(f2) + + return self.nodes + + def get_connected_components(self) -> List[Set[str]]: + """Find connected components in the dependency graph.""" + visited = set() + components = [] + + def dfs(node_path: str, component: Set[str]): + if node_path in visited: + return + visited.add(node_path) + component.add(node_path) + + node = self.nodes.get(node_path) + if not node: + return + + for related_path in node.related_files: + if related_path in self.nodes: + dfs(related_path, component) + + for path in self.nodes: + if path not in visited: + component: Set[str] = set() + dfs(path, component) + if component: + components.append(component) + + return components + + def get_smart_batches( + self, + file_groups: List[Any], + workspace: str, + project: str, + branches: List[str], + max_batch_size: int = 7, + min_batch_size: int = 3, + enrichment_data: Any = None + ) -> List[List[Dict[str, Any]]]: + """ + Create intelligent batches that keep related files together. + + Strategy: + 1. If enrichment_data is available, use pre-computed relationships from Java + 2. Otherwise, query RAG to discover file relationships via tree-sitter metadata + 3. Find connected components (files that are related) + 4. Batch files within components together + 5. For large components, split by priority while keeping related files together + + Args: + file_groups: List of FileGroup objects with files + workspace: Repository workspace/owner + project: Repository slug + branches: Branch names for context + max_batch_size: Maximum files per batch + min_batch_size: Minimum files per batch + enrichment_data: Optional PrEnrichmentDataDto from Java with pre-computed relationships + """ + # Use enrichment data if available, otherwise fall back to RAG + if enrichment_data and hasattr(enrichment_data, 'has_data') and enrichment_data.has_data(): + logger.info("Using pre-computed enrichment data for dependency graph") + self.build_graph_from_enrichment(file_groups, enrichment_data) + else: + self.build_graph_from_rag(file_groups, workspace, project, branches) + components = self.get_connected_components() + + logger.info( + f"Dependency analysis: {len(self.nodes)} files, " + f"{len(components)} connected components, " + f"{len(self.relationships)} relationships" + ) + + file_priority_map = {} + file_info_map = {} + for group in file_groups: + for f in group.files: + file_priority_map[f.path] = group.priority + file_info_map[f.path] = f + + batches = [] + processed_files = set() + priority_order = ['CRITICAL', 'HIGH', 'MEDIUM', 'LOW'] + + def component_sort_key(comp): + max_priority = min( + priority_order.index(file_priority_map.get(f, 'LOW')) + for f in comp + ) + return (-len(comp), max_priority) + + for component in sorted(components, key=component_sort_key): + if all(f in processed_files for f in component): + continue + + component_files = [f for f in component if f not in processed_files] + if not component_files: + continue + + component_files_sorted = sorted( + component_files, + key=lambda f: ( + -self.nodes[f].relationship_strength, + priority_order.index(file_priority_map.get(f, 'LOW')), + f + ) + ) + + current_batch = [] + for file_path in component_files_sorted: + file_info = file_info_map.get(file_path) + if not file_info: + continue + + node = self.nodes[file_path] + current_batch.append({ + "file": file_info, + "priority": file_priority_map.get(file_path, 'MEDIUM'), + "has_relationships": len(node.related_files) > 0, + "relationship_strength": node.relationship_strength, + "related_in_batch": [ + r for r in node.related_files + if r in {b['file'].path for b in current_batch} + ] + }) + processed_files.add(file_path) + + if len(current_batch) >= max_batch_size: + batches.append(current_batch) + current_batch = [] + + if current_batch: + batches.append(current_batch) + + # Handle orphan files + orphan_files = [] + for group in file_groups: + for f in group.files: + if f.path not in processed_files: + orphan_files.append({ + "file": f, + "priority": group.priority, + "has_relationships": False, + "relationship_strength": 0.0, + "related_in_batch": [] + }) + processed_files.add(f.path) + + if orphan_files: + orphan_files_sorted = sorted( + orphan_files, + key=lambda x: (priority_order.index(x['priority']), x['file'].path) + ) + for i in range(0, len(orphan_files_sorted), max_batch_size): + batches.append(orphan_files_sorted[i:i + max_batch_size]) + + batches = self._merge_small_batches(batches, min_batch_size, max_batch_size) + + logger.info(f"Smart batching created {len(batches)} batches from {len(self.nodes)} files") + for i, batch in enumerate(batches): + paths = [b['file'].path for b in batch] + rel_count = sum(1 for b in batch if b.get('has_relationships')) + logger.debug(f"Batch {i+1}: {len(batch)} files ({rel_count} with relationships): {paths}") + + return batches + + def _merge_small_batches( + self, + batches: List[List[Dict[str, Any]]], + min_size: int, + max_size: int + ) -> List[List[Dict[str, Any]]]: + """Merge small batches if they have the same priority.""" + if not batches: + return batches + + priority_batches: Dict[str, List[List[Dict[str, Any]]]] = defaultdict(list) + for batch in batches: + if not batch: + continue + priorities = [b['priority'] for b in batch] + dominant = max(set(priorities), key=priorities.count) + priority_batches[dominant].append(batch) + + merged = [] + for priority, p_batches in priority_batches.items(): + current_merged = [] + for batch in p_batches: + if len(current_merged) + len(batch) <= max_size: + current_merged.extend(batch) + else: + if current_merged: + merged.append(current_merged) + current_merged = batch[:] + if current_merged: + merged.append(current_merged) + + return merged + + def get_relationship_summary(self) -> Dict[str, Any]: + """Get a summary of discovered relationships.""" + relationship_types = defaultdict(int) + for rel in self.relationships: + relationship_types[rel.relationship_type] += 1 + + return { + "total_files": len(self.nodes), + "total_relationships": len(self.relationships), + "relationship_types": dict(relationship_types), + "files_with_relationships": sum( + 1 for node in self.nodes.values() + if len(node.related_files) > 0 + ), + "avg_relationships_per_file": ( + len(self.relationships) * 2 / len(self.nodes) + if self.nodes else 0 + ) + } + + +def create_smart_batches( + file_groups: List[Any], + workspace: str, + project: str, + branches: List[str], + rag_client: Optional["RAGClient"] = None, + max_batch_size: int = 7, + enrichment_data: Any = None +) -> List[List[Dict[str, Any]]]: + """ + Convenience function to create smart batches from file groups. + + Args: + file_groups: List of FileGroup objects with files + workspace: Repository workspace/owner + project: Repository slug + branches: Branch names for context + rag_client: Optional RAG client for relationship discovery + max_batch_size: Maximum files per batch + enrichment_data: Optional PrEnrichmentDataDto with pre-computed relationships from Java + """ + builder = DependencyGraphBuilder(rag_client=rag_client) + return builder.get_smart_batches( + file_groups, + workspace, + project, + branches, + max_batch_size, + enrichment_data=enrichment_data + ) diff --git a/python-ecosystem/mcp-client/utils/diff_processor.py b/python-ecosystem/mcp-client/utils/diff_processor.py index 913c2f04..17616bfe 100644 --- a/python-ecosystem/mcp-client/utils/diff_processor.py +++ b/python-ecosystem/mcp-client/utils/diff_processor.py @@ -50,7 +50,8 @@ class DiffFile: old_path: Optional[str] = None # For renamed files additions: int = 0 deletions: int = 0 - content: str = "" + content: str = "" # Diff content (unified diff format) + full_content: Optional[str] = None # Full file content (populated separately if needed) hunks: List[str] = field(default_factory=list) is_binary: bool = False is_skipped: bool = False diff --git a/python-ecosystem/mcp-client/utils/prompts/prompt_builder.py b/python-ecosystem/mcp-client/utils/prompts/prompt_builder.py index b3cad5b1..7883d90d 100644 --- a/python-ecosystem/mcp-client/utils/prompts/prompt_builder.py +++ b/python-ecosystem/mcp-client/utils/prompts/prompt_builder.py @@ -1,6 +1,6 @@ from typing import Any, Dict, List, Optional import json -from model.models import IssueDTO +from model.dtos import IssueDTO from utils.prompts.prompt_constants import ( ADDITIONAL_INSTRUCTIONS, BRANCH_REVIEW_PROMPT_TEMPLATE, @@ -73,7 +73,8 @@ def build_stage_1_batch_prompt( project_rules: str = "", rag_context: str = "", is_incremental: bool = False, - previous_issues: str = "" + previous_issues: str = "", + all_pr_files: List[str] = None # All files in this PR for cross-file awareness ) -> str: """ Build prompt for Stage 1: Batch File Review. @@ -103,6 +104,20 @@ def build_stage_1_batch_prompt( This is a follow-up review after the PR was updated with new commits. The diff above shows ONLY the changes since the last review - focus on these NEW changes. For any previous issues listed below, check if they are RESOLVED in the new changes. +""" + + # Add PR-wide file list for cross-batch awareness + pr_files_context = "" + if all_pr_files: + current_batch_files = [f['path'] for f in files] + other_files = [fp for fp in all_pr_files if fp not in current_batch_files] + if other_files: + pr_files_context = f""" +## OTHER FILES IN THIS PR (for cross-file awareness) +This PR also modifies these files (reviewed in other batches): +{chr(10).join('- ' + fp for fp in other_files[:20])} +{'... and ' + str(len(other_files) - 20) + ' more files' if len(other_files) > 20 else ''} +Consider potential interactions with these files when reviewing. """ return STAGE_1_BATCH_PROMPT_TEMPLATE.format( @@ -111,7 +126,8 @@ def build_stage_1_batch_prompt( files_context=files_context, rag_context=rag_context or "(No additional codebase context available)", incremental_instructions=incremental_instructions, - previous_issues=previous_issues + previous_issues=previous_issues, + pr_files_context=pr_files_context ) @staticmethod diff --git a/python-ecosystem/mcp-client/utils/prompts/prompt_constants.py b/python-ecosystem/mcp-client/utils/prompts/prompt_constants.py index 64d09c20..f9d79c70 100644 --- a/python-ecosystem/mcp-client/utils/prompts/prompt_constants.py +++ b/python-ecosystem/mcp-client/utils/prompts/prompt_constants.py @@ -275,7 +275,7 @@ Create a prioritized review plan in this JSON format: {{ - "analysis_summary": "2-sentence overview of PR scope and risk level", + "analysis_summary": "overview of PR scope and risk level", "file_groups": [ {{ "group_id": "GROUP_A_SECURITY", @@ -343,18 +343,26 @@ - Missing variable declarations - they may exist outside the diff context - Missing function definitions - the function may be defined elsewhere in the file - Missing class properties - they may be declared outside the visible changes +- Security issues in code that is not visible in the diff of RAG context ONLY report issues that you can VERIFY from the visible diff content. If you suspect an issue but cannot confirm it from the diff, DO NOT report it. When in doubt, assume the code is correct - the developer can see the full file, you cannot. {incremental_instructions} +{pr_files_context} PROJECT RULES: {project_rules} CODEBASE CONTEXT (from RAG): {rag_context} +IMPORTANT: When referencing codebase context in your analysis: +- ALWAYS cite the actual file path (e.g., "as seen in `src/service/UserService.java`") +- NEVER reference context by number (e.g., DO NOT say "Related Code #1" or "chunk #3") +- Quote relevant code snippets when needed to support your analysis +- The numbered headers are for your reference only, not for output + {previous_issues} SUGGESTED_FIX_DIFF_FORMAT: @@ -369,7 +377,7 @@ Review each file below independently. For each file, produce a review result. Use the CODEBASE CONTEXT above to understand how the changed code integrates with existing patterns, dependencies, and architectural decisions. -If previous issue fixed in a current version, mark it as resolved. +If a previous issue (from PREVIOUS ISSUES section) is fixed in the current version, include it with isResolved=true and preserve its original id. INPUT FILES: Priority: {priority} @@ -385,14 +393,15 @@ "analysis_summary": "Summary of findings for file 1", "issues": [ {{ + "id": "original-issue-id-if-from-previous-issues", "severity": "HIGH|MEDIUM|LOW|INFO", "category": "SECURITY|PERFORMANCE|CODE_QUALITY|BUG_RISK|STYLE|DOCUMENTATION|BEST_PRACTICES|ERROR_HANDLING|TESTING|ARCHITECTURE", "file": "path/to/file1", "line": "42", - "reason": "Detailed explanation of the issue", + "reason": "Detailed explanation of the issue (or resolution reason if isResolved=true)", "suggestedFixDescription": "Clear description of how to fix the issue", "suggestedFixDiff": "Unified diff showing exact code changes (MUST follow SUGGESTED_FIX_DIFF_FORMAT)", - "isResolved": false|true + "isResolved": false }} ], "confidence": "HIGH|MEDIUM|LOW|INFO", @@ -409,12 +418,14 @@ - Return exactly one review object per input file. - Match file paths exactly. - Skip style nits. +- For PREVIOUS ISSUES that are now RESOLVED: set "isResolved": true (boolean, not string) and PRESERVE the original id field. +- The "isResolved" field MUST be a JSON boolean: true or false, NOT a string "true" or "false". - suggestedFixDiff MUST be a valid unified diff string if a fix is proposed. """ STAGE_2_CROSS_FILE_PROMPT_TEMPLATE = """SYSTEM ROLE: You are a staff architect reviewing this PR for systemic risks. -Focus on: data flow, authorization patterns, consistency, database integrity, service boundaries. +Focus on: data flow, authorization patterns, consistency, service boundaries. At temperature 0.1, you will be conservative—that is correct. Flag even low-confidence concerns. Return structured JSON. @@ -436,7 +447,7 @@ Architecture Reference {architecture_context} -Database Migrations in This PR +Migration Files in This PR {migrations} Output Format @@ -465,15 +476,6 @@ "severity": "HIGH" }} ], - "immutability_enforcement": {{ - "rule": "Analysis results immutable after status=FINAL", - "check_pass": true, - "evidence": "..." - }}, - "database_integrity": {{ - "concerns": ["FK constraints", "cascade deletes"], - "findings": [] - }}, "pr_recommendation": "PASS|PASS_WITH_WARNINGS|FAIL", "confidence": "HIGH|MEDIUM|LOW|INFO" }} diff --git a/python-ecosystem/mcp-client/utils/response_parser.py b/python-ecosystem/mcp-client/utils/response_parser.py index 9336a7ee..68490e2f 100644 --- a/python-ecosystem/mcp-client/utils/response_parser.py +++ b/python-ecosystem/mcp-client/utils/response_parser.py @@ -1,7 +1,10 @@ import json +import logging import re from typing import Any, Dict, List, Union, Optional +logger = logging.getLogger(__name__) + class ResponseParser: """Parser class for extracting and structuring AI responses.""" @@ -12,7 +15,7 @@ class ResponseParser: # Valid issue fields - others will be removed VALID_ISSUE_FIELDS = { - 'id', 'severity', 'category', 'file', 'line', 'reason', + 'id', 'issueId', 'severity', 'category', 'file', 'line', 'reason', 'suggestedFixDescription', 'suggestedFixDiff', 'isResolved' } @@ -66,7 +69,7 @@ def _normalize_diff(diff_value: Any) -> Optional[str]: def _clean_issue(issue: Dict[str, Any]) -> Dict[str, Any]: """ Clean and normalize a single issue object. - - Removes unexpected fields (like 'id') + - Normalizes issueId to id (for consistent field naming) - Normalizes suggestedFixDiff format - Normalizes severity and category @@ -81,7 +84,15 @@ def _clean_issue(issue: Dict[str, Any]) -> Dict[str, Any]: cleaned = {} + # Normalize issueId to id (AI may use either) + if 'issueId' in issue and 'id' not in issue: + issue['id'] = issue['issueId'] + for key, value in issue.items(): + # Skip issueId since we normalized it to id above + if key == 'issueId': + continue + # Skip fields not in valid set if key not in ResponseParser.VALID_ISSUE_FIELDS: continue @@ -111,10 +122,15 @@ def _clean_issue(issue: Dict[str, Any]) -> Dict[str, Any]: # Ensure isResolved is boolean if key == 'isResolved': + original_value = value if isinstance(value, str): value = value.lower() == 'true' elif not isinstance(value, bool): value = False + # Log resolved issues for debugging + if value: + issue_id = issue.get('id') or issue.get('issueId', 'unknown') + logger.info(f"Issue {issue_id} marked as isResolved=True (original: {original_value})") # Ensure id is a string when present (preserve mapping to DB ids) if key == 'id': @@ -689,21 +705,14 @@ def create_error_response(error_message: str, exception_str: str = "") -> Dict[s exception_str: Optional exception details Returns: - Structured error response dictionary with issues as list + Structured error response dictionary marked as error (no fake issues) """ full_message = f"{error_message}: {exception_str}" if exception_str else error_message return { + "status": "error", "comment": full_message, - "issues": [ - { - "severity": "HIGH", - "category": "ERROR_HANDLING", - "file": "system", - "line": "0", - "reason": full_message, - "suggestedFixDescription": "Check system configuration and connectivity", - "isResolved": False - } - ] + "issues": [], # Don't create fake issues for errors - let Java handle error state properly + "error": True, + "error_message": full_message } \ No newline at end of file diff --git a/python-ecosystem/rag-pipeline/Dockerfile b/python-ecosystem/rag-pipeline/Dockerfile index 46f619c5..5537e2e7 100644 --- a/python-ecosystem/rag-pipeline/Dockerfile +++ b/python-ecosystem/rag-pipeline/Dockerfile @@ -34,6 +34,13 @@ ENV LLAMA_INDEX_CACHE_DIR=/tmp/.llama_index # Allow concurrent indexing by running multiple Uvicorn workers ENV UVICORN_WORKERS=4 +# CPU Threading Optimization (can be overridden via .env) +# These control parallelism for numerical operations +ENV OMP_NUM_THREADS=4 +ENV MKL_NUM_THREADS=4 +ENV OPENBLAS_NUM_THREADS=4 +ENV NUMEXPR_NUM_THREADS=4 + EXPOSE 8001 CMD ["python", "main.py"] diff --git a/python-ecosystem/rag-pipeline/main.py b/python-ecosystem/rag-pipeline/main.py index 7b3e146f..3e18e67c 100644 --- a/python-ecosystem/rag-pipeline/main.py +++ b/python-ecosystem/rag-pipeline/main.py @@ -12,37 +12,57 @@ format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) -load_dotenv() +load_dotenv(interpolate=False) # Validate critical environment variables before starting def validate_environment(): """Validate that required environment variables are set""" + embedding_provider = os.environ.get("EMBEDDING_PROVIDER", "ollama").lower() openrouter_key = os.environ.get("OPENROUTER_API_KEY", "") logger.info("=" * 60) logger.info("RAG Pipeline Starting - Environment Check") logger.info("=" * 60) - logger.info(f"QDRANT_URL: {os.getenv('QDRANT_URL', 'NOT SET')}") + logger.info(f"QDRANT_URL: {os.getenv('QDRANT_URL', 'http://localhost:6333')}") logger.info(f"QDRANT_COLLECTION_PREFIX: {os.getenv('QDRANT_COLLECTION_PREFIX', 'codecrow')}") - logger.info(f"OPENROUTER_MODEL: {os.getenv('OPENROUTER_MODEL', 'openai/text-embedding-3-small')}") + logger.info(f"EMBEDDING_PROVIDER: {embedding_provider}") - if not openrouter_key or openrouter_key.strip() == "": - logger.error("=" * 60) - logger.error("CRITICAL ERROR: OPENROUTER_API_KEY not set!") - logger.error("=" * 60) - logger.error("The OPENROUTER_API_KEY environment variable is required") - logger.error("but was not found or is empty.") - logger.error("") - logger.error("To fix this:") - logger.error("1. Set the environment variable:") - logger.error(" export OPENROUTER_API_KEY='sk-or-v1-...'") - logger.error("2. Or add it to docker-compose.yml") - logger.error("3. Or create a .env file with: OPENROUTER_API_KEY=sk-or-v1-...") - logger.error("=" * 60) - sys.exit(1) + if embedding_provider == "ollama": + ollama_url = os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434') + ollama_model = os.getenv('OLLAMA_EMBEDDING_MODEL', 'qwen3-embedding:0.6b') + logger.info(f"OLLAMA_BASE_URL: {ollama_url}") + logger.info(f"OLLAMA_EMBEDDING_MODEL: {ollama_model}") + logger.info("=" * 60) + logger.info("Using Ollama for local embeddings ✓") + logger.info("Make sure Ollama is running: ollama serve") + logger.info(f"And model is pulled: ollama pull {ollama_model}") + logger.info("=" * 60) + elif embedding_provider == "openrouter": + logger.info(f"OPENROUTER_MODEL: {os.getenv('OPENROUTER_MODEL', 'qwen/qwen3-embedding-8b')}") + + if not openrouter_key or openrouter_key.strip() == "": + logger.error("=" * 60) + logger.error("CRITICAL ERROR: OPENROUTER_API_KEY not set!") + logger.error("=" * 60) + logger.error("The OPENROUTER_API_KEY environment variable is required") + logger.error("when EMBEDDING_PROVIDER=openrouter but was not found or is empty.") + logger.error("") + logger.error("To fix this:") + logger.error("1. Set the environment variable:") + logger.error(" export OPENROUTER_API_KEY='sk-or-v1-...'") + logger.error("2. Or add it to docker-compose.yml") + logger.error("3. Or create a .env file with: OPENROUTER_API_KEY=sk-or-v1-...") + logger.error("4. Or switch to local embeddings: EMBEDDING_PROVIDER=ollama") + logger.error("=" * 60) + sys.exit(1) + + logger.info(f"OPENROUTER_API_KEY: {openrouter_key[:15]}...{openrouter_key[-4:]} ✓") + logger.info("=" * 60) + logger.info("Using OpenRouter for cloud embeddings ✓") + logger.info("=" * 60) + else: + logger.warning(f"Unknown EMBEDDING_PROVIDER '{embedding_provider}', defaulting to 'ollama'") - logger.info(f"OPENROUTER_API_KEY: {openrouter_key[:15]}...{openrouter_key[-4:]} ✓") - logger.info("=" * 60) logger.info("Environment validation passed ✓") logger.info("Using Qdrant for vector storage") logger.info("=" * 60) diff --git a/python-ecosystem/rag-pipeline/scripts/examples.py b/python-ecosystem/rag-pipeline/scripts/examples.py deleted file mode 100644 index f2b0146a..00000000 --- a/python-ecosystem/rag-pipeline/scripts/examples.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Example usage of the RAG pipeline -""" - -from pathlib import Path -from rag_pipeline.models.config import RAGConfig -from rag_pipeline.core.index_manager import RAGIndexManager -from rag_pipeline.services.query_service import RAGQueryService -from rag_pipeline.services.webhook_integration import WebhookIntegration - - -def example_index_repository(): - """Example: Index a repository""" - print("=== Example: Index Repository ===\n") - - config = RAGConfig() - index_manager = RAGIndexManager(config) - - stats = index_manager.index_repository( - repo_path="/path/to/your/repo", - workspace="codecrow", - project="demo", - branch="main", - commit="abc123" - ) - - print(f"Indexed {stats.document_count} documents") - print(f"Created {stats.chunk_count} chunks") - print(f"Namespace: {stats.namespace}") - print() - - -def example_incremental_update(): - """Example: Incrementally update files""" - print("=== Example: Incremental Update ===\n") - - config = RAGConfig() - index_manager = RAGIndexManager(config) - - changed_files = ["src/main.py", "src/utils.py"] - - stats = index_manager.update_files( - file_paths=changed_files, - repo_base="/path/to/your/repo", - workspace="codecrow", - project="demo", - branch="feature/new", - commit="def456" - ) - - print(f"Updated {len(changed_files)} files") - print(f"Total chunks: {stats.chunk_count}") - print() - - -def example_semantic_search(): - """Example: Perform semantic search""" - print("=== Example: Semantic Search ===\n") - - config = RAGConfig() - query_service = RAGQueryService(config) - - results = query_service.semantic_search( - query="authentication implementation", - workspace="codecrow", - project="demo", - branch="main", - top_k=5 - ) - - for i, result in enumerate(results, 1): - print(f"Result {i} (Score: {result['score']:.3f})") - print(f"File: {result['metadata']['path']}") - print(f"Language: {result['metadata']['language']}") - print(f"Preview: {result['text'][:100]}...") - print() - - -def example_pr_context(): - """Example: Get context for PR review""" - print("=== Example: PR Context ===\n") - - config = RAGConfig() - query_service = RAGQueryService(config) - - context = query_service.get_context_for_pr( - workspace="codecrow", - project="demo", - branch="main", - changed_files=["src/auth.py", "src/user.py"], - pr_description="Implemented new JWT authentication", - top_k=10 - ) - - print("Context for PR review:") - print(context) - print() - - -def example_webhook_integration(): - """Example: Webhook integration""" - print("=== Example: Webhook Integration ===\n") - - config = RAGConfig() - webhook = WebhookIntegration(config) - - # Handle PR created - result = webhook.handle_pr_created( - workspace="codecrow", - project="demo", - branch="feature/new", - commit="abc123", - repo_clone_path="/tmp/repo_clone", - pr_description="New feature implementation" - ) - - print(f"PR Created: {result['success']}") - if result['success']: - print(f"Message: {result['message']}") - print(f"Stats: {result['stats']}") - - # Handle PR updated - result = webhook.handle_pr_updated( - workspace="codecrow", - project="demo", - branch="feature/new", - commit="def456", - changed_files=["src/main.py", "src/utils.py"], - repo_clone_path="/tmp/repo_clone" - ) - - print(f"\nPR Updated: {result['success']}") - if result['success']: - print(f"Files updated: {result['files_updated']}") - print(f"Files deleted: {result['files_deleted']}") - - # Get context for analysis - context = webhook.get_pr_analysis_context( - workspace="codecrow", - project="demo", - branch="feature/new", - changed_files=["src/main.py"], - pr_description="Updated main logic" - ) - - print(f"\nAnalysis Context Length: {len(context)} characters") - print() - - -def example_list_indices(): - """Example: List all indices""" - print("=== Example: List Indices ===\n") - - config = RAGConfig() - index_manager = RAGIndexManager(config) - - indices = index_manager.list_indices() - - print(f"Found {len(indices)} indices:") - for idx in indices: - print(f"- {idx.namespace}") - print(f" Workspace: {idx.workspace}, Project: {idx.project}, Branch: {idx.branch}") - print(f" Documents: {idx.document_count}, Chunks: {idx.chunk_count}") - print(f" Last Updated: {idx.last_updated}") - print() - - -def example_delete_index(): - """Example: Delete an index""" - print("=== Example: Delete Index ===\n") - - config = RAGConfig() - index_manager = RAGIndexManager(config) - - index_manager.delete_index( - workspace="codecrow", - project="demo", - branch="old-feature" - ) - - print("Index deleted successfully") - print() - - -if __name__ == "__main__": - print("CodeCrow RAG Pipeline - Examples\n") - print("Note: Update the paths and credentials before running\n") - - # Uncomment the examples you want to run: - - # example_index_repository() - # example_incremental_update() - # example_semantic_search() - # example_pr_context() - # example_webhook_integration() - # example_list_indices() - # example_delete_index() - - print("Done!") - diff --git a/python-ecosystem/rag-pipeline/scripts/quick-start.sh b/python-ecosystem/rag-pipeline/scripts/quick-start.sh deleted file mode 100755 index 11890e82..00000000 --- a/python-ecosystem/rag-pipeline/scripts/quick-start.sh +++ /dev/null @@ -1,137 +0,0 @@ -#!/bin/bash - -set -e - -# Change to project root -cd "$(dirname "$0")/.." - -echo "=========================================" -echo "CodeCrow RAG Pipeline - Quick Start" -echo "=========================================" -echo "" - -# Check if .env exists -if [ ! -f .env ]; then - echo "❌ .env file not found!" - echo "" - echo "Please create .env file with:" - echo " OPENROUTER_API_KEY=your-key-here" - echo "" - echo "You can copy from .env.sample:" - echo " cp .env.sample .env" - echo " nano .env" - exit 1 -fi - -# Check if OPENROUTER_API_KEY is set -source .env -if [ -z "$OPENROUTER_API_KEY" ] || [ "$OPENROUTER_API_KEY" = "your_openrouter_api_key_here" ]; then - echo "❌ OPENROUTER_API_KEY not configured in .env" - echo "" - echo "Please update .env with your OpenRouter API key" - exit 1 -fi - -echo "✅ Configuration found" -echo "" - -# Check if Docker is running -if ! docker info > /dev/null 2>&1; then - echo "❌ Docker is not running" - echo "" - echo "Please start Docker and try again" - exit 1 -fi - -echo "✅ Docker is running" -echo "" - -# Check if running from root directory -if [ ! -f "../docker-compose.yml" ]; then - echo "❌ docker-compose.yml not found" - echo "" - echo "Please run this script from the codecrow-rag-pipeline directory" - echo "Or use the root-level start-rag-pipeline.sh script" - exit 1 -fi - -echo "Starting RAG Pipeline services..." -echo "" - -# Start MongoDB -echo "Starting MongoDB..." -cd .. && docker-compose up -d mongodb && cd codecrow-rag-pipeline - -# Wait for MongoDB to be ready -echo "Waiting for MongoDB to be ready..." -sleep 10 - -# Check MongoDB health -if docker exec codecrow-mongodb mongosh --eval "db.adminCommand('ping')" > /dev/null 2>&1; then - echo "✅ MongoDB is ready" -else - echo "⚠️ MongoDB may still be starting up" -fi - -echo "" - -# Build and start RAG pipeline -echo "Building and starting RAG Pipeline..." -cd .. && docker-compose up -d --build rag-pipeline && cd codecrow-rag-pipeline - -# Wait for service to start -echo "Waiting for RAG Pipeline to start..." -sleep 5 - -# Check health -echo "" -echo "Checking RAG Pipeline health..." -for i in {1..10}; do - if curl -s http://localhost:8001/health > /dev/null 2>&1; then - echo "✅ RAG Pipeline is healthy" - break - fi - - if [ $i -eq 10 ]; then - echo "⚠️ RAG Pipeline may still be starting" - echo " Check logs: docker logs codecrow-rag-pipeline" - else - sleep 2 - fi -done - -echo "" -echo "=========================================" -echo "Services Started Successfully!" -echo "=========================================" -echo "" -echo "📊 Service Status:" -cd .. && docker-compose ps mongodb rag-pipeline && cd codecrow-rag-pipeline -echo "" - -echo "🔗 Endpoints:" -echo " - RAG API: http://localhost:8001" -echo " - Health Check: http://localhost:8001/health" -echo " - API Docs: http://localhost:8001/docs" -echo " - MongoDB: mongodb://localhost:27017" -echo "" - -echo "📝 Quick Test:" -echo " curl http://localhost:8001/health" -echo "" - -echo "📚 Documentation:" -echo " - README.md - Module overview" -echo " - INTEGRATION_GUIDE.md - How to integrate" -echo " - DEPLOYMENT.md - Deployment guide" -echo " - IMPLEMENTATION_SUMMARY.md - What was built" -echo "" - -echo "🛠️ Useful Commands:" -echo " - View logs: docker logs -f codecrow-rag-pipeline" -echo " - Stop services: docker-compose stop mongodb rag-pipeline" -echo " - Restart: docker-compose restart rag-pipeline" -echo "" - -echo "✅ Setup complete! RAG Pipeline is ready to use." - diff --git a/python-ecosystem/rag-pipeline/scripts/setup.sh b/python-ecosystem/rag-pipeline/scripts/setup.sh deleted file mode 100755 index fddc82e2..00000000 --- a/python-ecosystem/rag-pipeline/scripts/setup.sh +++ /dev/null @@ -1,43 +0,0 @@ -#!/bin/bash - -set -e - -echo "Setting up CodeCrow RAG Pipeline..." - -# Check if Python 3.11+ is available -python_version=$(python3 --version 2>&1 | awk '{print $2}' | cut -d. -f1,2) -echo "Python version: $python_version" - -# Create virtual environment -echo "Creating virtual environment..." -python3 -m venv .venv - -# Activate virtual environment -echo "Activating virtual environment..." -source .venv/bin/activate - -# Upgrade pip -echo "Upgrading pip..." -pip install --upgrade pip - -# Install requirements -echo "Installing requirements..." -pip install -r requirements.txt - -# Create .env from sample if not exists -if [ ! -f .env ]; then - echo "Creating .env from .env.sample..." - cp .env.sample .env - echo "Please update .env with your configuration" -fi - -echo "" -echo "Setup complete!" -echo "" -echo "To activate the virtual environment, run:" -echo " source .venv/bin/activate" -echo "" -echo "To start the API server, run:" -echo " python main.py" -echo "" -echo "Don't forget to configure MongoDB and OpenAI API key in .env" diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/__init__.py index 3dc56254..dbafd454 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/__init__.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/__init__.py @@ -2,27 +2,25 @@ CodeCrow RAG Pipeline A RAG (Retrieval-Augmented Generation) pipeline for code repositories. -Provides indexing and querying capabilities for code using LlamaIndex and MongoDB. +Provides indexing and querying capabilities for code using LlamaIndex, Tree-sitter and Qdrant. """ __version__ = "1.0.0" -from .models.config import RAGConfig, DocumentMetadata, IndexStats +from .models.config import RAGConfig, IndexStats from .core.index_manager import RAGIndexManager from .services.query_service import RAGQueryService from .core.loader import DocumentLoader -from .core.chunking import CodeAwareSplitter, FunctionAwareSplitter +from .core.splitter import ASTCodeSplitter from .utils.utils import make_namespace, detect_language_from_path __all__ = [ "RAGConfig", - "DocumentMetadata", "IndexStats", "RAGIndexManager", "RAGQueryService", "DocumentLoader", - "CodeAwareSplitter", - "FunctionAwareSplitter", + "ASTCodeSplitter", "make_namespace", "detect_language_from_path", ] diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/api/api.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/api/api.py index 0b97486d..35190448 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/api/api.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/api/api.py @@ -1,8 +1,12 @@ import logging import gc -from typing import List, Optional +import os +import uuid +from typing import Dict, List, Optional from fastapi import FastAPI, HTTPException, BackgroundTasks -from pydantic import BaseModel +from pydantic import BaseModel, field_validator +from llama_index.core.schema import TextNode +from qdrant_client.models import PointStruct from ..models.config import RAGConfig, IndexStats from ..core.index_manager import RAGIndexManager @@ -13,11 +17,29 @@ app = FastAPI(title="CodeCrow RAG API", version="2.0.0") +# Service-to-service auth +from .middleware import ServiceSecretMiddleware +app.add_middleware(ServiceSecretMiddleware) + config = RAGConfig() index_manager = RAGIndexManager(config) query_service = RAGQueryService(config) +# Allowed base directory for all repo paths (set via env or default to /tmp) +_ALLOWED_REPO_ROOT = os.environ.get("ALLOWED_REPO_ROOT", "/tmp") + + +def _validate_repo_path(path: str) -> str: + """Validate that a repo path is within the allowed root and contains no traversal.""" + resolved = os.path.realpath(path) + if not resolved.startswith(os.path.realpath(_ALLOWED_REPO_ROOT)): + raise ValueError( + f"Path must be under {_ALLOWED_REPO_ROOT}, got: {path}" + ) + return path + + class IndexRequest(BaseModel): repo_path: str workspace: str @@ -26,6 +48,11 @@ class IndexRequest(BaseModel): commit: str exclude_patterns: Optional[List[str]] = None + @field_validator("repo_path") + @classmethod + def validate_repo_path(cls, v: str) -> str: + return _validate_repo_path(v) + class UpdateFilesRequest(BaseModel): file_paths: List[str] @@ -35,6 +62,11 @@ class UpdateFilesRequest(BaseModel): branch: str commit: str + @field_validator("repo_base") + @classmethod + def validate_repo_base(cls, v: str) -> str: + return _validate_repo_path(v) + class DeleteFilesRequest(BaseModel): file_paths: List[str] @@ -65,6 +97,26 @@ class PRContextRequest(BaseModel): enable_priority_reranking: Optional[bool] = True min_relevance_score: Optional[float] = 0.7 deleted_files: Optional[List[str]] = [] # Files deleted in target branch + # NEW: PR-specific hybrid query mode + pr_number: Optional[int] = None # If set, enables hybrid query with PR data priority + all_pr_changed_files: Optional[List[str]] = [] # All files in PR (for exclusion from branch query) + + +class DeterministicContextRequest(BaseModel): + """Request for deterministic metadata-based context retrieval. + + TWO-STEP process leveraging tree-sitter metadata: + 1. Get chunks for the changed file_paths + 2. Extract semantic_names/imports/extends from those chunks + 3. Find related definitions using extracted identifiers + + NO language-specific parsing needed - tree-sitter already did it during indexing! + """ + workspace: str + project: str + branches: List[str] # Branches to search (e.g., ['release/1.29', 'master']) + file_paths: List[str] # Changed file paths from diff + limit_per_file: Optional[int] = 10 # Max chunks per file class DeleteBranchRequest(BaseModel): @@ -83,6 +135,160 @@ def health(): return {"status": "healthy"} +# ============================================================================= +# PARSE ENDPOINTS (AST metadata extraction without indexing) +# ============================================================================= + +class ParseFileRequest(BaseModel): + """Request to parse a single file and extract AST metadata.""" + path: str + content: str + language: Optional[str] = None # Auto-detected if not provided + + +class ParseBatchRequest(BaseModel): + """Request to parse multiple files in batch.""" + files: List[ParseFileRequest] + + +class ParsedFileMetadata(BaseModel): + """AST metadata extracted from a file.""" + path: str + language: Optional[str] = None + imports: List[str] = [] + extends: List[str] = [] + implements: List[str] = [] + semantic_names: List[str] = [] # Function/class/method names + parent_class: Optional[str] = None + namespace: Optional[str] = None + calls: List[str] = [] # Called functions/methods + success: bool = True + error: Optional[str] = None + + +@app.post("/parse", response_model=ParsedFileMetadata) +def parse_file(request: ParseFileRequest): + """ + Parse a single file and extract AST metadata WITHOUT indexing. + + Returns tree-sitter extracted metadata: + - imports: Import statements + - extends: Parent classes/interfaces + - implements: Implemented interfaces + - semantic_names: Function/class/method names defined + - namespace: Package/namespace + - calls: Called functions/methods + + Used by Java pipeline-agent to build dependency graph. + """ + try: + from ..core.splitter import ASTCodeSplitter + from ..core.splitter.languages import get_language_from_path, EXTENSION_TO_LANGUAGE + + # Detect language + language = request.language + if not language: + lang_enum = get_language_from_path(request.path) + language = lang_enum.value if lang_enum else None + + if not language: + # Try to infer from extension + ext = '.' + request.path.rsplit('.', 1)[-1] if '.' in request.path else '' + language = EXTENSION_TO_LANGUAGE.get(ext, {}).get('name') + + splitter = ASTCodeSplitter( + max_chunk_size=50000, # Large to avoid splitting + enrich_embedding_text=False + ) + + # Create a minimal document for parsing + from llama_index.core.schema import Document as LlamaDocument + doc = LlamaDocument(text=request.content, metadata={'path': request.path}) + + # Parse and extract chunks + nodes = splitter.split_documents([doc]) + + # Aggregate metadata from all chunks + imports = set() + extends = set() + implements = set() + semantic_names = set() + calls = set() + namespace = None + parent_classes = set() + + for node in nodes: + meta = node.metadata + + if meta.get('imports'): + imports.update(meta['imports']) + if meta.get('extends'): + extends.update(meta['extends']) + if meta.get('implements'): + implements.update(meta['implements']) + if meta.get('semantic_names'): + semantic_names.update(meta['semantic_names']) + if meta.get('calls'): + calls.update(meta['calls']) + if meta.get('namespace') and not namespace: + namespace = meta['namespace'] + if meta.get('parent_class'): + parent_classes.add(meta['parent_class']) + + # Primary parent class (first one found) + parent_class = list(parent_classes)[0] if parent_classes else None + + return ParsedFileMetadata( + path=request.path, + language=language, + imports=sorted(list(imports)), + extends=sorted(list(extends)), + implements=sorted(list(implements)), + semantic_names=sorted(list(semantic_names)), + parent_class=parent_class, + namespace=namespace, + calls=sorted(list(calls)), + success=True + ) + + except Exception as e: + logger.warning(f"Error parsing file {request.path}: {e}") + return ParsedFileMetadata( + path=request.path, + success=False, + error=str(e) + ) + + +@app.post("/parse/batch") +def parse_files_batch(request: ParseBatchRequest): + """ + Parse multiple files and extract AST metadata in batch. + + Returns list of ParsedFileMetadata for each file. + Continues processing even if individual files fail. + """ + results = [] + + for file_req in request.files: + result = parse_file(file_req) + results.append(result) + + successful = sum(1 for r in results if r.success) + failed = len(results) - successful + + logger.info(f"Batch parse: {successful} successful, {failed} failed out of {len(results)} files") + + return { + "results": results, + "summary": { + "total": len(results), + "successful": successful, + "failed": failed + } + } + + @app.get("/limits") def get_limits(): """Get current RAG indexing limits (for free plan info)""" @@ -99,6 +305,11 @@ class EstimateRequest(BaseModel): repo_path: str exclude_patterns: Optional[List[str]] = None + @field_validator("repo_path") + @classmethod + def validate_repo_path(cls, v: str) -> str: + return _validate_repo_path(v) + class EstimateResponse(BaseModel): file_count: int @@ -365,15 +576,19 @@ def semantic_search(request: QueryRequest): @app.post("/query/pr-context") def get_pr_context(request: PRContextRequest): """ - Get context for PR review with multi-branch support. + Get context for PR review with multi-branch support and optional PR-specific hybrid mode. - Queries both target branch and base branch to preserve cross-file relationships. - Results are deduplicated with target branch taking priority. + When pr_number is provided, uses HYBRID query: + 1. Query PR-indexed chunks (pr=true, pr_number=X) - these are the actual changed files + 2. Query branch data, excluding files that are in the PR (to get unchanged dependencies) + 3. Merge results with PR data taking priority Args: branch: Target branch (PR source branch) base_branch: Base branch (PR target, e.g., 'main'). Auto-detected if not provided. deleted_files: Files deleted in target branch (excluded from results) + pr_number: If set, enables hybrid query with PR data priority + all_pr_changed_files: Files to exclude from branch query (PR files already indexed separately) """ try: # If branch is not provided, return empty context @@ -392,6 +607,21 @@ def get_pr_context(request: PRContextRequest): } } + pr_results = [] + + # HYBRID MODE: Query PR-indexed data first if pr_number is provided + if request.pr_number: + pr_results = _query_pr_indexed_data( + workspace=request.workspace, + project=request.project, + pr_number=request.pr_number, + query_texts=request.diff_snippets or [], + pr_title=request.pr_title, + top_k=request.top_k or 15 + ) + logger.info(f"Hybrid mode: Found {len(pr_results)} PR-specific chunks for PR #{request.pr_number}") + + # Get branch context (with optional exclusion of PR files) context = query_service.get_context_for_pr( workspace=request.workspace, project=request.project, @@ -404,16 +634,42 @@ def get_pr_context(request: PRContextRequest): enable_priority_reranking=request.enable_priority_reranking, min_relevance_score=request.min_relevance_score, base_branch=request.base_branch, - deleted_files=request.deleted_files or [] + deleted_files=request.deleted_files or [], + # Pass PR files to exclude if in hybrid mode (guard against None) + exclude_pr_files=(request.all_pr_changed_files or []) if request.pr_number else [] ) + # Merge PR results with branch results (PR first, then branch) + if pr_results: + pr_paths = set() + merged_code = [] + + # PR results first (highest priority - fresh data) + # Include ALL chunks from PR (multiple chunks per file allowed) + for pr_chunk in pr_results: + merged_code.append(pr_chunk) + path = pr_chunk.get("path", "") + if path: + pr_paths.add(path) + + # Then branch results (excluding files already covered by PR) + for branch_chunk in context.get("relevant_code", []): + path = branch_chunk.get("path", "") + if path not in pr_paths: + merged_code.append(branch_chunk) + + context["relevant_code"] = merged_code + context["_pr_chunks_count"] = len(pr_results) + # Add metadata about processing context["_metadata"] = { "priority_reranking_enabled": request.enable_priority_reranking, "min_relevance_score": request.min_relevance_score, "changed_files_count": len(request.changed_files), "result_count": len(context.get("relevant_code", [])), - "branches_searched": context.get("_branches_searched", [request.branch]) + "branches_searched": context.get("_branches_searched", [request.branch]), + "hybrid_mode": request.pr_number is not None, + "pr_number": request.pr_number } return {"context": context} @@ -422,6 +678,104 @@ def get_pr_context(request: PRContextRequest): raise HTTPException(status_code=500, detail=str(e)) +def _query_pr_indexed_data( + workspace: str, + project: str, + pr_number: int, + query_texts: List[str], + pr_title: Optional[str], + top_k: int = 15 +) -> List[Dict]: + """ + Query PR-indexed chunks from the main collection. + + Filters by pr=true and pr_number to get only PR-specific data. + """ + from qdrant_client.models import Filter, FieldCondition, MatchValue + + collection_name = index_manager._get_project_collection_name(workspace, project) + + if not index_manager._collection_manager.collection_exists(collection_name): + return [] + + # Build query from snippets and title + query_parts = [] + if pr_title: + query_parts.append(pr_title) + if query_texts: + query_parts.extend(query_texts[:5]) # Limit snippets + + if not query_parts: + # If no query, just get all PR chunks + query_text = f"code changes for PR {pr_number}" + else: + query_text = " ".join(query_parts)[:1000] + + try: + # Generate embedding for query + query_embedding = index_manager.embed_model.get_text_embedding(query_text) + + # Search with PR filter + results = index_manager.qdrant_client.search( + collection_name=collection_name, + query_vector=query_embedding, + query_filter=Filter( + must=[ + FieldCondition(key="pr", match=MatchValue(value=True)), + FieldCondition(key="pr_number", match=MatchValue(value=pr_number)) + ] + ), + limit=top_k, + with_payload=True + ) + + formatted = [] + for r in results: + formatted.append({ + "path": r.payload.get("path", "unknown"), + "content": r.payload.get("text", ""), + "score": r.score, + "semantic_name": r.payload.get("semantic_name", ""), + "semantic_type": r.payload.get("semantic_type", ""), + "branch": r.payload.get("pr_branch", ""), + "_source": "pr_indexed" + }) + + return formatted + + except Exception as e: + logger.warning(f"Error querying PR-indexed data: {e}") + return [] + + +@app.post("/query/deterministic") +def get_deterministic_context(request: DeterministicContextRequest): + """ + Get context using DETERMINISTIC metadata-based retrieval. + + Two-step process: + 1. Get chunks for changed file_paths + 2. Extract semantic_names/imports/extends from those chunks (tree-sitter metadata!) + 3. Find related definitions using extracted identifiers + + No language-specific parsing needed - tree-sitter already did it during indexing. + Predictable: same input = same output. + """ + try: + context = query_service.get_deterministic_context( + workspace=request.workspace, + project=request.project, + branches=request.branches, + file_paths=request.file_paths, + limit_per_file=request.limit_per_file or 10 + ) + + return {"context": context} + except Exception as e: + logger.error(f"Error getting deterministic context: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + # ============================================================================= # BRANCH MANAGEMENT ENDPOINTS # ============================================================================= @@ -546,6 +900,184 @@ def get_memory_usage(): raise HTTPException(status_code=500, detail=str(e)) +# ============================================================================= +# PR-SPECIFIC RAG ENDPOINTS (for PR file indexing with metadata) +# ============================================================================= + +class PRFileInfo(BaseModel): + """Info about a single PR file.""" + path: str + content: str # Full file content (not just diff) + change_type: str # ADDED, MODIFIED, DELETED + + +class PRIndexRequest(BaseModel): + """Request to index PR files into main collection with PR metadata.""" + workspace: str + project: str + pr_number: int + branch: str # Source branch + files: List[PRFileInfo] + + +@app.post("/index/pr-files") +def index_pr_files(request: PRIndexRequest): + """ + Index PR files into the main collection with PR-specific metadata. + + Files are indexed with metadata: + - pr: true + - pr_number: + - pr_branch: + + This allows hybrid queries that prioritize PR data over branch data. + Existing PR points for the same pr_number are deleted first. + """ + try: + from datetime import datetime, timezone + from llama_index.core import Document as LlamaDocument + + collection_name = index_manager._get_project_collection_name( + request.workspace, request.project + ) + + # Ensure collection exists + index_manager._ensure_collection_exists(collection_name) + + # Delete existing points for this PR first (handles re-analysis) + try: + from qdrant_client.models import Filter, FieldCondition, MatchValue + index_manager.qdrant_client.delete( + collection_name=collection_name, + points_selector=Filter( + must=[ + FieldCondition(key="pr_number", match=MatchValue(value=request.pr_number)) + ] + ) + ) + logger.info(f"Deleted existing PR points for PR #{request.pr_number}") + except Exception as e: + logger.warning(f"Error deleting existing PR points: {e}") + + # Convert files to LlamaIndex documents + documents = [] + for file_info in request.files: + if not file_info.content or not file_info.content.strip(): + continue + if file_info.change_type == "DELETED": + continue # Don't index deleted files + + doc = LlamaDocument( + text=file_info.content, + metadata={ + "path": file_info.path, + "change_type": file_info.change_type, + } + ) + documents.append(doc) + + if not documents: + return { + "status": "skipped", + "message": "No files to index", + "chunks_indexed": 0 + } + + # Split documents into chunks + chunks = index_manager.splitter.split_documents(documents) + + # Add PR metadata to all chunks + for chunk in chunks: + chunk.metadata["pr"] = True + chunk.metadata["pr_number"] = request.pr_number + chunk.metadata["pr_branch"] = request.branch + chunk.metadata["workspace"] = request.workspace + chunk.metadata["project"] = request.project + chunk.metadata["branch"] = request.branch # For compatibility + chunk.metadata["indexed_at"] = datetime.now(timezone.utc).isoformat() + + # Embed and upsert using point_ops (with PR-specific IDs) + chunk_data = [] + chunks_by_file = {} + for chunk in chunks: + path = chunk.metadata.get("path", str(uuid.uuid4())) + if path not in chunks_by_file: + chunks_by_file[path] = [] + chunks_by_file[path].append(chunk) + + for path, file_chunks in chunks_by_file.items(): + for chunk_index, chunk in enumerate(file_chunks): + # Use PR-specific ID format to avoid collision with branch data + key = f"pr:{request.pr_number}:{request.workspace}:{request.project}:{path}:{chunk_index}" + point_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, key)) + chunk_data.append((point_id, chunk)) + + # Embed and create points + points = index_manager._point_ops.embed_and_create_points(chunk_data) + + # Upsert to collection + successful, failed = index_manager._point_ops.upsert_points(collection_name, points) + + logger.info(f"Indexed PR #{request.pr_number}: {successful} chunks from {len(documents)} files") + + return { + "status": "indexed", + "pr_number": request.pr_number, + "files_processed": len(documents), + "chunks_indexed": successful, + "chunks_failed": failed + } + + except ValueError as e: + logger.warning(f"Invalid request for PR indexing: {e}") + raise HTTPException(status_code=400, detail=str(e)) + except Exception as e: + logger.error(f"Internal error indexing PR files: {e}") + raise HTTPException(status_code=500, detail="Internal indexing error") + + +@app.delete("/index/pr-files/{workspace}/{project}/{pr_number}") +def delete_pr_files(workspace: str, project: str, pr_number: int): + """ + Delete all indexed points for a specific PR. + + Called after analysis completes to clean up PR-specific data. + """ + try: + from qdrant_client.models import Filter, FieldCondition, MatchValue + + collection_name = index_manager._get_project_collection_name(workspace, project) + + # Check if collection exists + if not index_manager._collection_manager.collection_exists(collection_name): + return { + "status": "skipped", + "message": f"Collection does not exist" + } + + # Delete points with matching pr_number + result = index_manager.qdrant_client.delete( + collection_name=collection_name, + points_selector=Filter( + must=[ + FieldCondition(key="pr_number", match=MatchValue(value=pr_number)) + ] + ) + ) + + logger.info(f"Deleted PR #{pr_number} points from {collection_name}") + + return { + "status": "deleted", + "pr_number": pr_number, + "collection": collection_name + } + + except Exception as e: + logger.error(f"Error deleting PR files: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8001) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/api/middleware.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/api/middleware.py new file mode 100644 index 00000000..96bef76c --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/api/middleware.py @@ -0,0 +1,56 @@ +""" +Shared-secret authentication middleware. + +Validates that internal service-to-service requests carry +the correct X-Service-Secret header matching the SERVICE_SECRET env var. +""" +import os +import logging + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse + +logger = logging.getLogger(__name__) + +# Paths that skip auth (health checks, readiness probes) +_PUBLIC_PATHS = frozenset({"/health", "/docs", "/openapi.json", "/redoc"}) + + +class ServiceSecretMiddleware(BaseHTTPMiddleware): + """Reject requests that don't carry a valid shared service secret.""" + + def __init__(self, app, secret: str | None = None): + super().__init__(app) + self.secret = secret or os.environ.get("SERVICE_SECRET", "") + if self.secret: + logger.info("ServiceSecretMiddleware: secret configured (length=%d)", len(self.secret)) + else: + logger.warning("ServiceSecretMiddleware: no secret configured — auth disabled") + + async def dispatch(self, request: Request, call_next): + # Skip auth for health/doc endpoints + if request.url.path in _PUBLIC_PATHS: + return await call_next(request) + + # If no secret is configured, allow all (dev mode) + if not self.secret: + return await call_next(request) + + provided = request.headers.get("x-service-secret", "") + if provided != self.secret: + logger.warning( + "Unauthorized request to %s from %s — " + "provided_len=%d expected_len=%d match=%s", + request.url.path, + request.client.host if request.client else "unknown", + len(provided), + len(self.secret), + provided == self.secret, + ) + return JSONResponse( + status_code=401, + content={"detail": "Missing or invalid service secret"}, + ) + + return await call_next(request) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/__init__.py index dcb908ff..43cff58b 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/__init__.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/__init__.py @@ -1,17 +1,10 @@ """Core functionality for indexing and document processing""" __all__ = [ "DocumentLoader", - "CodeAwareSplitter", - "FunctionAwareSplitter", - "SemanticCodeSplitter", "ASTCodeSplitter", "RAGIndexManager" ] from .index_manager import RAGIndexManager -from .chunking import CodeAwareSplitter, FunctionAwareSplitter -from .semantic_splitter import SemanticCodeSplitter -from .ast_splitter import ASTCodeSplitter -from .loader import DocumentLoader - - +from .splitter import ASTCodeSplitter +from .loader import DocumentLoader \ No newline at end of file diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/ast_splitter.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/ast_splitter.py deleted file mode 100644 index 29e5369b..00000000 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/ast_splitter.py +++ /dev/null @@ -1,1396 +0,0 @@ -""" -AST-based Code Splitter using Tree-sitter for accurate code parsing. - -This module provides true AST-aware code chunking that: -1. Uses Tree-sitter for accurate AST parsing (15+ languages) -2. Splits code into semantic units (classes, functions, methods) -3. Uses RecursiveCharacterTextSplitter for oversized chunks (large methods) -4. Enriches metadata for better RAG retrieval -5. Maintains parent context ("breadcrumbs") for nested structures -6. Uses deterministic IDs for Qdrant deduplication - -Key benefits over regex-based splitting: -- Accurate function/class boundary detection -- Language-aware parsing for 15+ languages -- Better metadata: content_type, language, semantic_names, parent_class -- Handles edge cases (nested functions, decorators, etc.) -- Deterministic chunk IDs prevent duplicates on re-indexing -""" - -import re -import hashlib -import logging -from typing import List, Dict, Any, Optional, Set -from pathlib import Path -from dataclasses import dataclass, field -from enum import Enum - -from langchain_text_splitters import RecursiveCharacterTextSplitter, Language -from llama_index.core.schema import Document as LlamaDocument, TextNode - -logger = logging.getLogger(__name__) - - -class ContentType(Enum): - """Content type as determined by AST parsing""" - FUNCTIONS_CLASSES = "functions_classes" # Full function/class definition - SIMPLIFIED_CODE = "simplified_code" # Remaining code with placeholders - FALLBACK = "fallback" # Non-AST parsed content - OVERSIZED_SPLIT = "oversized_split" # Large chunk split by RecursiveCharacterTextSplitter - - -# Map file extensions to LangChain Language enum -EXTENSION_TO_LANGUAGE: Dict[str, Language] = { - # Python - '.py': Language.PYTHON, - '.pyw': Language.PYTHON, - '.pyi': Language.PYTHON, - - # Java/JVM - '.java': Language.JAVA, - '.kt': Language.KOTLIN, - '.kts': Language.KOTLIN, - '.scala': Language.SCALA, - - # JavaScript/TypeScript - '.js': Language.JS, - '.jsx': Language.JS, - '.mjs': Language.JS, - '.cjs': Language.JS, - '.ts': Language.TS, - '.tsx': Language.TS, - - # Systems languages - '.go': Language.GO, - '.rs': Language.RUST, - '.c': Language.C, - '.h': Language.C, - '.cpp': Language.CPP, - '.cc': Language.CPP, - '.cxx': Language.CPP, - '.hpp': Language.CPP, - '.hxx': Language.CPP, - '.cs': Language.CSHARP, - - # Web/Scripting - '.php': Language.PHP, - '.phtml': Language.PHP, # PHP template files (Magento, Zend, etc.) - '.php3': Language.PHP, - '.php4': Language.PHP, - '.php5': Language.PHP, - '.phps': Language.PHP, - '.inc': Language.PHP, # PHP include files - '.rb': Language.RUBY, - '.erb': Language.RUBY, # Ruby template files - '.lua': Language.LUA, - '.pl': Language.PERL, - '.pm': Language.PERL, - '.swift': Language.SWIFT, - - # Markup/Config - '.md': Language.MARKDOWN, - '.markdown': Language.MARKDOWN, - '.html': Language.HTML, - '.htm': Language.HTML, - '.rst': Language.RST, - '.tex': Language.LATEX, - '.proto': Language.PROTO, - '.sol': Language.SOL, - '.hs': Language.HASKELL, - '.cob': Language.COBOL, - '.cbl': Language.COBOL, - '.xml': Language.HTML, # Use HTML splitter for XML -} - -# Languages that support full AST parsing via tree-sitter -AST_SUPPORTED_LANGUAGES = { - Language.PYTHON, Language.JAVA, Language.KOTLIN, Language.JS, Language.TS, - Language.GO, Language.RUST, Language.C, Language.CPP, Language.CSHARP, - Language.PHP, Language.RUBY, Language.SCALA, Language.LUA, Language.PERL, - Language.SWIFT, Language.HASKELL, Language.COBOL -} - -# Tree-sitter language name mapping (tree-sitter-languages uses these names) -LANGUAGE_TO_TREESITTER: Dict[Language, str] = { - Language.PYTHON: 'python', - Language.JAVA: 'java', - Language.KOTLIN: 'kotlin', - Language.JS: 'javascript', - Language.TS: 'typescript', - Language.GO: 'go', - Language.RUST: 'rust', - Language.C: 'c', - Language.CPP: 'cpp', - Language.CSHARP: 'c_sharp', - Language.PHP: 'php', - Language.RUBY: 'ruby', - Language.SCALA: 'scala', - Language.LUA: 'lua', - Language.PERL: 'perl', - Language.SWIFT: 'swift', - Language.HASKELL: 'haskell', -} - -# Node types that represent semantic CHUNKING units (classes, functions) -# NOTE: imports, namespace, inheritance are now extracted DYNAMICALLY from AST -# by pattern matching on node type names - no hardcoded mappings needed! -SEMANTIC_NODE_TYPES: Dict[str, Dict[str, List[str]]] = { - 'python': { - 'class': ['class_definition'], - 'function': ['function_definition', 'async_function_definition'], - }, - 'java': { - 'class': ['class_declaration', 'interface_declaration', 'enum_declaration'], - 'function': ['method_declaration', 'constructor_declaration'], - }, - 'javascript': { - 'class': ['class_declaration'], - 'function': ['function_declaration', 'method_definition', 'arrow_function', 'generator_function_declaration'], - }, - 'typescript': { - 'class': ['class_declaration', 'interface_declaration'], - 'function': ['function_declaration', 'method_definition', 'arrow_function'], - }, - 'go': { - 'class': ['type_declaration'], # structs, interfaces - 'function': ['function_declaration', 'method_declaration'], - }, - 'rust': { - 'class': ['struct_item', 'impl_item', 'trait_item', 'enum_item'], - 'function': ['function_item'], - }, - 'c_sharp': { - 'class': ['class_declaration', 'interface_declaration', 'struct_declaration'], - 'function': ['method_declaration', 'constructor_declaration'], - }, - 'kotlin': { - 'class': ['class_declaration', 'object_declaration', 'interface_declaration'], - 'function': ['function_declaration'], - }, - 'php': { - 'class': ['class_declaration', 'interface_declaration', 'trait_declaration'], - 'function': ['function_definition', 'method_declaration'], - }, - 'ruby': { - 'class': ['class', 'module'], - 'function': ['method', 'singleton_method'], - }, - 'cpp': { - 'class': ['class_specifier', 'struct_specifier'], - 'function': ['function_definition'], - }, - 'c': { - 'class': ['struct_specifier'], - 'function': ['function_definition'], - }, - 'scala': { - 'class': ['class_definition', 'object_definition', 'trait_definition'], - 'function': ['function_definition'], - }, -} - -# Metadata extraction patterns (fallback when AST doesn't provide names) -METADATA_PATTERNS = { - 'python': { - 'class': re.compile(r'^class\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'^(?:async\s+)?def\s+(\w+)\s*\(', re.MULTILINE), - }, - 'java': { - 'class': re.compile(r'(?:public\s+|private\s+|protected\s+)?(?:abstract\s+|final\s+)?class\s+(\w+)', re.MULTILINE), - 'interface': re.compile(r'(?:public\s+)?interface\s+(\w+)', re.MULTILINE), - 'method': re.compile(r'(?:public|private|protected)\s+(?:static\s+)?[\w<>,\s]+\s+(\w+)\s*\(', re.MULTILINE), - }, - 'javascript': { - 'class': re.compile(r'(?:export\s+)?(?:default\s+)?class\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:export\s+)?(?:async\s+)?function\s*\*?\s*(\w+)\s*\(', re.MULTILINE), - }, - 'typescript': { - 'class': re.compile(r'(?:export\s+)?(?:default\s+)?class\s+(\w+)', re.MULTILINE), - 'interface': re.compile(r'(?:export\s+)?interface\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:export\s+)?(?:async\s+)?function\s*\*?\s*(\w+)\s*\(', re.MULTILINE), - }, - 'go': { - 'function': re.compile(r'^func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(', re.MULTILINE), - 'struct': re.compile(r'^type\s+(\w+)\s+struct\s*\{', re.MULTILINE), - }, - 'rust': { - 'function': re.compile(r'^(?:pub\s+)?(?:async\s+)?fn\s+(\w+)', re.MULTILINE), - 'struct': re.compile(r'^(?:pub\s+)?struct\s+(\w+)', re.MULTILINE), - }, - 'c_sharp': { - 'class': re.compile(r'(?:public\s+|private\s+|internal\s+)?(?:abstract\s+|sealed\s+)?class\s+(\w+)', re.MULTILINE), - 'method': re.compile(r'(?:public|private|protected|internal)\s+(?:static\s+)?[\w<>,\s]+\s+(\w+)\s*\(', re.MULTILINE), - }, - 'kotlin': { - 'class': re.compile(r'(?:data\s+|sealed\s+|open\s+)?class\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:fun|suspend\s+fun)\s+(\w+)\s*\(', re.MULTILINE), - }, - 'php': { - 'class': re.compile(r'(?:abstract\s+|final\s+)?class\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:public|private|protected|static|\s)*function\s+(\w+)\s*\(', re.MULTILINE), - }, -} - -# Patterns for extracting class inheritance, interfaces, and imports -CLASS_INHERITANCE_PATTERNS = { - 'php': { - 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w\\]+)', re.MULTILINE), - 'implements': re.compile(r'class\s+\w+(?:\s+extends\s+[\w\\]+)?\s+implements\s+([\w\\,\s]+)', re.MULTILINE), - 'use': re.compile(r'^use\s+([\w\\]+)(?:\s+as\s+\w+)?;', re.MULTILINE), - 'namespace': re.compile(r'^namespace\s+([\w\\]+);', re.MULTILINE), - 'type_hint': re.compile(r'@var\s+(\\?[\w\\|]+)', re.MULTILINE), - # PHTML template type hints: /** @var \Namespace\Class $variable */ - 'template_type': re.compile(r'/\*\*\s*@var\s+([\w\\]+)\s+\$\w+\s*\*/', re.MULTILINE), - # Variable type hints in PHPDoc: @param Type $name, @return Type - 'phpdoc_types': re.compile(r'@(?:param|return|throws)\s+([\w\\|]+)', re.MULTILINE), - }, - 'java': { - 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), - 'implements': re.compile(r'class\s+\w+(?:\s+extends\s+[\w.]+)?\s+implements\s+([\w.,\s]+)', re.MULTILINE), - 'import': re.compile(r'^import\s+([\w.]+(?:\.\*)?);', re.MULTILINE), - 'package': re.compile(r'^package\s+([\w.]+);', re.MULTILINE), - }, - 'kotlin': { - 'extends': re.compile(r'class\s+\w+\s*:\s*([\w.]+)(?:\([^)]*\))?', re.MULTILINE), - 'import': re.compile(r'^import\s+([\w.]+(?:\.\*)?)', re.MULTILINE), - 'package': re.compile(r'^package\s+([\w.]+)', re.MULTILINE), - }, - 'python': { - 'extends': re.compile(r'class\s+\w+\s*\(\s*([\w.,\s]+)\s*\)\s*:', re.MULTILINE), - 'import': re.compile(r'^(?:from\s+([\w.]+)\s+)?import\s+([\w.,\s*]+)', re.MULTILINE), - }, - 'typescript': { - 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), - 'implements': re.compile(r'class\s+\w+(?:\s+extends\s+[\w.]+)?\s+implements\s+([\w.,\s]+)', re.MULTILINE), - 'import': re.compile(r'^import\s+(?:[\w{},\s*]+\s+from\s+)?["\']([^"\']+)["\'];?', re.MULTILINE), - }, - 'javascript': { - 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), - 'import': re.compile(r'^import\s+(?:[\w{},\s*]+\s+from\s+)?["\']([^"\']+)["\'];?', re.MULTILINE), - 'require': re.compile(r'require\s*\(\s*["\']([^"\']+)["\']\s*\)', re.MULTILINE), - }, - 'c_sharp': { - 'extends': re.compile(r'class\s+\w+\s*:\s*([\w.]+)', re.MULTILINE), - 'implements': re.compile(r'class\s+\w+\s*:\s*(?:[\w.]+\s*,\s*)*([\w.,\s]+)', re.MULTILINE), - 'using': re.compile(r'^using\s+([\w.]+);', re.MULTILINE), - 'namespace': re.compile(r'^namespace\s+([\w.]+)', re.MULTILINE), - }, - 'go': { - 'import': re.compile(r'^import\s+(?:\(\s*)?"([^"]+)"', re.MULTILINE), - 'package': re.compile(r'^package\s+(\w+)', re.MULTILINE), - }, - 'rust': { - 'use': re.compile(r'^use\s+([\w:]+(?:::\{[^}]+\})?);', re.MULTILINE), - 'impl_for': re.compile(r'impl\s+(?:<[^>]+>\s+)?([\w:]+)\s+for\s+([\w:]+)', re.MULTILINE), - }, - 'scala': { - 'extends': re.compile(r'(?:class|object|trait)\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), - 'with': re.compile(r'with\s+([\w.]+)', re.MULTILINE), - 'import': re.compile(r'^import\s+([\w._{}]+)', re.MULTILINE), - 'package': re.compile(r'^package\s+([\w.]+)', re.MULTILINE), - }, -} - - -@dataclass -class ASTChunk: - """Represents a chunk of code from AST parsing""" - content: str - content_type: ContentType - language: str - path: str - semantic_names: List[str] = field(default_factory=list) - parent_context: List[str] = field(default_factory=list) # Breadcrumb: ["MyClass", "inner_method"] - docstring: Optional[str] = None - signature: Optional[str] = None - start_line: int = 0 - end_line: int = 0 - node_type: Optional[str] = None - class_metadata: Dict[str, Any] = field(default_factory=dict) # extends, implements from AST - file_metadata: Dict[str, Any] = field(default_factory=dict) # imports, namespace from AST - - -def generate_deterministic_id(path: str, content: str, chunk_index: int = 0) -> str: - """ - Generate a deterministic ID for a chunk based on file path and content. - - This ensures the same code chunk always gets the same ID, preventing - duplicates in Qdrant during re-indexing. - - Args: - path: File path - content: Chunk content - chunk_index: Index of chunk within file (for disambiguation) - - Returns: - Deterministic hex ID string - """ - # Use path + content hash + index for uniqueness - hash_input = f"{path}:{chunk_index}:{content[:500]}" # First 500 chars for efficiency - return hashlib.sha256(hash_input.encode('utf-8')).hexdigest()[:32] - - -def compute_file_hash(content: str) -> str: - """Compute hash of file content for change detection""" - return hashlib.sha256(content.encode('utf-8')).hexdigest() - - -class ASTCodeSplitter: - """ - AST-based code splitter using Tree-sitter for accurate parsing. - - Features: - - True AST parsing via tree-sitter for accurate code structure detection - - Splits code into semantic units (classes, functions, methods) - - Maintains parent context (breadcrumbs) for nested structures - - Falls back to RecursiveCharacterTextSplitter for oversized chunks - - Uses deterministic IDs for Qdrant deduplication - - Enriches metadata for improved RAG retrieval - - Usage: - splitter = ASTCodeSplitter(max_chunk_size=2000) - nodes = splitter.split_documents(documents) - """ - - DEFAULT_MAX_CHUNK_SIZE = 2000 - DEFAULT_MIN_CHUNK_SIZE = 100 - DEFAULT_CHUNK_OVERLAP = 200 - DEFAULT_PARSER_THRESHOLD = 10 # Minimum lines for AST parsing - - def __init__( - self, - max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE, - min_chunk_size: int = DEFAULT_MIN_CHUNK_SIZE, - chunk_overlap: int = DEFAULT_CHUNK_OVERLAP, - parser_threshold: int = DEFAULT_PARSER_THRESHOLD - ): - """ - Initialize AST code splitter. - - Args: - max_chunk_size: Maximum characters per chunk. Larger chunks are split. - min_chunk_size: Minimum characters for a valid chunk. - chunk_overlap: Overlap between chunks when splitting oversized content. - parser_threshold: Minimum lines for AST parsing (smaller files use fallback). - """ - self.max_chunk_size = max_chunk_size - self.min_chunk_size = min_chunk_size - self.chunk_overlap = chunk_overlap - self.parser_threshold = parser_threshold - - # Cache text splitters for oversized chunks - self._splitter_cache: Dict[Language, RecursiveCharacterTextSplitter] = {} - - # Default text splitter for unknown languages - self._default_splitter = RecursiveCharacterTextSplitter( - chunk_size=max_chunk_size, - chunk_overlap=chunk_overlap, - length_function=len, - ) - - # Track if tree-sitter is available - self._tree_sitter_available: Optional[bool] = None - # Cache for language modules and parsers - self._language_cache: Dict[str, Any] = {} - - def _get_tree_sitter_language(self, lang_name: str): - """ - Get tree-sitter Language object for a language name. - Uses the new tree-sitter API with individual language packages. - - Note: Different packages have different APIs: - - Most use: module.language() - - PHP uses: module.language_php() - - TypeScript uses: module.language_typescript() - """ - if lang_name in self._language_cache: - return self._language_cache[lang_name] - - try: - from tree_sitter import Language - - # Map language names to their package modules and function names - # Format: (module_name, function_name or None for 'language') - lang_modules = { - 'python': ('tree_sitter_python', 'language'), - 'java': ('tree_sitter_java', 'language'), - 'javascript': ('tree_sitter_javascript', 'language'), - 'typescript': ('tree_sitter_typescript', 'language_typescript'), - 'go': ('tree_sitter_go', 'language'), - 'rust': ('tree_sitter_rust', 'language'), - 'c': ('tree_sitter_c', 'language'), - 'cpp': ('tree_sitter_cpp', 'language'), - 'c_sharp': ('tree_sitter_c_sharp', 'language'), - 'ruby': ('tree_sitter_ruby', 'language'), - 'php': ('tree_sitter_php', 'language_php'), - } - - lang_info = lang_modules.get(lang_name) - if not lang_info: - return None - - module_name, func_name = lang_info - - # Dynamic import of language module - import importlib - lang_module = importlib.import_module(module_name) - - # Get the language function - lang_func = getattr(lang_module, func_name, None) - if not lang_func: - logger.debug(f"Module {module_name} has no {func_name} function") - return None - - # Create Language object using the new API - language = Language(lang_func()) - self._language_cache[lang_name] = language - return language - - except Exception as e: - logger.debug(f"Could not load tree-sitter language '{lang_name}': {e}") - return None - - def _check_tree_sitter(self) -> bool: - """Check if tree-sitter is available""" - if self._tree_sitter_available is None: - try: - from tree_sitter import Parser, Language - import tree_sitter_python as tspython - - # Test with the new API - py_language = Language(tspython.language()) - parser = Parser(py_language) - parser.parse(b"def test(): pass") - - self._tree_sitter_available = True - logger.info("tree-sitter is available and working") - except ImportError as e: - logger.warning(f"tree-sitter not installed: {e}") - self._tree_sitter_available = False - except Exception as e: - logger.warning(f"tree-sitter error: {type(e).__name__}: {e}") - self._tree_sitter_available = False - return self._tree_sitter_available - - def _get_language_from_path(self, path: str) -> Optional[Language]: - """Determine Language enum from file path""" - ext = Path(path).suffix.lower() - return EXTENSION_TO_LANGUAGE.get(ext) - - def _get_treesitter_language(self, language: Language) -> Optional[str]: - """Get tree-sitter language name from Language enum""" - return LANGUAGE_TO_TREESITTER.get(language) - - def _get_text_splitter(self, language: Language) -> RecursiveCharacterTextSplitter: - """Get language-specific text splitter for oversized chunks""" - if language not in self._splitter_cache: - try: - self._splitter_cache[language] = RecursiveCharacterTextSplitter.from_language( - language=language, - chunk_size=self.max_chunk_size, - chunk_overlap=self.chunk_overlap, - ) - except Exception: - # Fallback if language not supported - self._splitter_cache[language] = self._default_splitter - return self._splitter_cache[language] - - def _parse_inheritance_clause(self, clause_text: str, language: str) -> List[str]: - """ - Parse inheritance clause from AST node text to extract class/interface names. - - Handles various formats: - - PHP: "extends ParentClass" or "implements Interface1, Interface2" - - Java: "extends Parent" or "implements I1, I2" - - Python: "(Parent1, Parent2)" - - TypeScript/JS: "extends Parent implements Interface" - """ - if not clause_text: - return [] - - # Clean up the clause - text = clause_text.strip() - - # Remove common keywords - for keyword in ['extends', 'implements', 'with', ':']: - text = text.replace(keyword, ' ') - - # Handle parentheses (Python style) - text = text.strip('()') - - # Split by comma and clean up - names = [] - for part in text.split(','): - name = part.strip() - # Remove any generic type parameters - if '<' in name: - name = name.split('<')[0].strip() - # Remove any constructor calls () - if '(' in name: - name = name.split('(')[0].strip() - if name and name not in ('', ' '): - names.append(name) - - return names - - def _parse_namespace(self, ns_text: str, language: str) -> Optional[str]: - """ - Extract namespace/package name from AST node text. - - Handles various formats: - - PHP: "namespace Vendor\\Package\\Module;" - - Java/Kotlin: "package com.example.app;" - - C#: "namespace MyNamespace { ... }" - """ - if not ns_text: - return None - - # Clean up - text = ns_text.strip() - - # Remove keywords and semicolons - for keyword in ['namespace', 'package']: - text = text.replace(keyword, ' ') - - text = text.strip().rstrip(';').rstrip('{').strip() - - return text if text else None - - def _parse_with_ast( - self, - text: str, - language: Language, - path: str - ) -> List[ASTChunk]: - """ - Parse code using AST via tree-sitter. - - Returns list of ASTChunk objects with content and metadata. - """ - if not self._check_tree_sitter(): - return [] - - ts_lang = self._get_treesitter_language(language) - if not ts_lang: - logger.debug(f"No tree-sitter mapping for {language}, using fallback") - return [] - - try: - from tree_sitter import Parser - - # Get Language object for this language - lang_obj = self._get_tree_sitter_language(ts_lang) - if not lang_obj: - logger.debug(f"tree-sitter language '{ts_lang}' not available") - return [] - - # Create parser with the language - parser = Parser(lang_obj) - tree = parser.parse(bytes(text, "utf8")) - - # Extract chunks with breadcrumb context - chunks = self._extract_ast_chunks_with_context( - tree.root_node, - text, - ts_lang, - path - ) - - return chunks - - except Exception as e: - logger.warning(f"AST parsing failed for {path}: {e}") - return [] - - def _extract_ast_chunks_with_context( - self, - root_node, - source_code: str, - language: str, - path: str - ) -> List[ASTChunk]: - """ - Extract function/class chunks from AST tree with parent context (breadcrumbs). - - This solves the "context loss" problem by tracking parent classes/modules - so that a method knows it belongs to a specific class. - - Also extracts file-level metadata dynamically from AST - no hardcoded mappings needed. - """ - chunks = [] - processed_ranges: Set[tuple] = set() # Track (start, end) to avoid duplicates - - # File-level metadata collected dynamically from AST - file_metadata: Dict[str, Any] = { - 'imports': [], - 'types_referenced': [], - } - - # Get node types for this language (only for chunking - class/function boundaries) - lang_node_types = SEMANTIC_NODE_TYPES.get(language, {}) - class_types = set(lang_node_types.get('class', [])) - function_types = set(lang_node_types.get('function', [])) - all_semantic_types = class_types | function_types - - def get_node_text(node) -> str: - """Get full text content of a node""" - return source_code[node.start_byte:node.end_byte] - - def extract_identifiers(node) -> List[str]: - """Recursively extract all identifier names from a node""" - identifiers = [] - if node.type in ('identifier', 'name', 'type_identifier', 'qualified_name', 'scoped_identifier'): - identifiers.append(get_node_text(node)) - for child in node.children: - identifiers.extend(extract_identifiers(child)) - return identifiers - - def extract_ast_metadata(node, metadata: Dict[str, Any]) -> None: - """ - Dynamically extract metadata from AST node based on node type patterns. - No hardcoded language mappings - uses common naming conventions in tree-sitter. - """ - node_type = node.type - node_text = get_node_text(node).strip() - - # === IMPORTS (pattern: *import*, *use*, *require*, *include*) === - if any(kw in node_type for kw in ('import', 'use_', 'require', 'include', 'using')): - if node_text and len(node_text) < 500: # Skip huge nodes - metadata['imports'].append(node_text) - return # Don't recurse into import children - - # === NAMESPACE/PACKAGE (pattern: *namespace*, *package*, *module*) === - if any(kw in node_type for kw in ('namespace', 'package', 'module_declaration')): - # Extract the name part - names = extract_identifiers(node) - if names: - metadata['namespace'] = names[0] if len(names) == 1 else '.'.join(names) - elif node_text: - # Fallback: parse from text - metadata['namespace'] = self._parse_namespace(node_text, language) - return - - # Recurse into children - for child in node.children: - extract_ast_metadata(child, metadata) - - def extract_class_metadata_from_ast(node) -> Dict[str, Any]: - """ - Dynamically extract class metadata (extends, implements) from AST. - Uses common tree-sitter naming patterns - no manual mapping needed. - """ - meta: Dict[str, Any] = {} - - def find_inheritance(n, depth=0): - """Recursively find inheritance-related nodes""" - node_type = n.type - - # === EXTENDS / SUPERCLASS (pattern: *super*, *base*, *extends*, *heritage*) === - if any(kw in node_type for kw in ('super', 'base_clause', 'extends', 'heritage', 'parent')): - names = extract_identifiers(n) - if names: - meta.setdefault('extends', []).extend(names) - meta['parent_types'] = meta.get('extends', []) - return # Found it, don't go deeper - - # === IMPLEMENTS / INTERFACES (pattern: *implement*, *interface_clause*, *conform*) === - if any(kw in node_type for kw in ('implement', 'interface_clause', 'conform', 'protocol')): - names = extract_identifiers(n) - if names: - meta.setdefault('implements', []).extend(names) - return - - # === TRAIT/MIXIN (pattern: *trait*, *mixin*, *with*) === - if any(kw in node_type for kw in ('trait', 'mixin', 'with_clause')): - names = extract_identifiers(n) - if names: - meta.setdefault('traits', []).extend(names) - return - - # === TYPE PARAMETERS / GENERICS === - if any(kw in node_type for kw in ('type_parameter', 'generic', 'type_argument')): - names = extract_identifiers(n) - if names: - meta.setdefault('type_params', []).extend(names) - return - - # Recurse but limit depth to avoid going too deep - if depth < 5: - for child in n.children: - find_inheritance(child, depth + 1) - - for child in node.children: - find_inheritance(child) - - # Deduplicate - for key in meta: - if isinstance(meta[key], list): - meta[key] = list(dict.fromkeys(meta[key])) # Preserve order, remove dupes - - return meta - - def get_node_name(node) -> Optional[str]: - """Extract name from a node (class/function name)""" - for child in node.children: - if child.type in ('identifier', 'name', 'type_identifier', 'property_identifier'): - return get_node_text(child) - return None - - def traverse(node, parent_context: List[str], depth: int = 0): - """ - Recursively traverse AST and extract semantic chunks with breadcrumbs. - - Args: - node: Current AST node - parent_context: List of parent class/function names (breadcrumb) - depth: Current depth in tree - """ - node_range = (node.start_byte, node.end_byte) - - # Check if this is a semantic unit - if node.type in all_semantic_types: - # Skip if already processed (nested in another chunk) - if node_range in processed_ranges: - return - - content = source_code[node.start_byte:node.end_byte] - - # Calculate line numbers - start_line = source_code[:node.start_byte].count('\n') + 1 - end_line = start_line + content.count('\n') - - # Get the name of this node - node_name = get_node_name(node) - - # Determine content type - is_class = node.type in class_types - - # Extract class inheritance metadata dynamically from AST - class_metadata = {} - if is_class: - class_metadata = extract_class_metadata_from_ast(node) - - chunk = ASTChunk( - content=content, - content_type=ContentType.FUNCTIONS_CLASSES, - language=language, - path=path, - semantic_names=[node_name] if node_name else [], - parent_context=list(parent_context), # Copy the breadcrumb - start_line=start_line, - end_line=end_line, - node_type=node.type, - class_metadata=class_metadata, - ) - - chunks.append(chunk) - processed_ranges.add(node_range) - - # If this is a class, traverse children with updated context - if is_class and node_name: - new_context = parent_context + [node_name] - for child in node.children: - traverse(child, new_context, depth + 1) - else: - # Continue traversing children with current context - for child in node.children: - traverse(child, parent_context, depth + 1) - - # First pass: extract file-level metadata (imports, namespace) from entire AST - extract_ast_metadata(root_node, file_metadata) - - # Second pass: extract semantic chunks (classes, functions) - traverse(root_node, []) - - # Clean up file_metadata - remove empty values - clean_file_metadata = {k: v for k, v in file_metadata.items() if v} - - # Create simplified code (skeleton with placeholders) - simplified = self._create_simplified_code(source_code, chunks, language) - if simplified and simplified.strip() and len(simplified.strip()) > 50: - chunks.append(ASTChunk( - content=simplified, - content_type=ContentType.SIMPLIFIED_CODE, - language=language, - path=path, - semantic_names=[], - parent_context=[], - start_line=1, - end_line=source_code.count('\n') + 1, - node_type='simplified', - file_metadata=clean_file_metadata, # Include imports/namespace from AST - )) - - # Also attach file_metadata to all chunks for enriched metadata - for chunk in chunks: - if not chunk.file_metadata: - chunk.file_metadata = clean_file_metadata - - return chunks - - def _create_simplified_code( - self, - source_code: str, - chunks: List[ASTChunk], - language: str - ) -> str: - """ - Create simplified code with placeholders for extracted chunks. - - This gives RAG context about the overall file structure without - including full function/class bodies. - - Example output: - # Code for: class MyClass: - # Code for: def my_function(): - if __name__ == "__main__": - main() - """ - if not chunks: - return source_code - - # Get chunks that are functions_classes type (not simplified) - semantic_chunks = [c for c in chunks if c.content_type == ContentType.FUNCTIONS_CLASSES] - - if not semantic_chunks: - return source_code - - # Sort by start position (reverse) to replace from end - sorted_chunks = sorted( - semantic_chunks, - key=lambda x: source_code.find(x.content), - reverse=True - ) - - result = source_code - - # Comment style by language - comment_prefix = { - 'python': '#', - 'javascript': '//', - 'typescript': '//', - 'java': '//', - 'kotlin': '//', - 'go': '//', - 'rust': '//', - 'c': '//', - 'cpp': '//', - 'c_sharp': '//', - 'php': '//', - 'ruby': '#', - 'lua': '--', - 'perl': '#', - 'scala': '//', - }.get(language, '//') - - for chunk in sorted_chunks: - # Find the position of this chunk in the source - pos = result.find(chunk.content) - if pos == -1: - continue - - # Extract first line for placeholder - first_line = chunk.content.split('\n')[0].strip() - # Truncate if too long - if len(first_line) > 60: - first_line = first_line[:60] + '...' - - # Add breadcrumb context to placeholder - breadcrumb = "" - if chunk.parent_context: - breadcrumb = f" (in {'.'.join(chunk.parent_context)})" - - placeholder = f"{comment_prefix} Code for: {first_line}{breadcrumb}\n" - - result = result[:pos] + placeholder + result[pos + len(chunk.content):] - - return result.strip() - - def _extract_metadata( - self, - chunk: ASTChunk, - base_metadata: Dict[str, Any] - ) -> Dict[str, Any]: - """Extract and enrich metadata from an AST chunk""" - metadata = dict(base_metadata) - - # Core AST metadata - metadata['content_type'] = chunk.content_type.value - metadata['node_type'] = chunk.node_type - - # Breadcrumb context (critical for RAG) - if chunk.parent_context: - metadata['parent_context'] = chunk.parent_context - metadata['parent_class'] = chunk.parent_context[-1] if chunk.parent_context else None - metadata['full_path'] = '.'.join(chunk.parent_context + chunk.semantic_names[:1]) - - # Semantic names - if chunk.semantic_names: - metadata['semantic_names'] = chunk.semantic_names[:10] - metadata['primary_name'] = chunk.semantic_names[0] - - # Line numbers - metadata['start_line'] = chunk.start_line - metadata['end_line'] = chunk.end_line - - # === Use AST-extracted metadata (from tree-sitter) === - - # File-level metadata from AST (imports, namespace, package) - if chunk.file_metadata: - if chunk.file_metadata.get('imports'): - metadata['imports'] = chunk.file_metadata['imports'][:20] - if chunk.file_metadata.get('namespace'): - metadata['namespace'] = chunk.file_metadata['namespace'] - if chunk.file_metadata.get('package'): - metadata['package'] = chunk.file_metadata['package'] - - # Class-level metadata from AST (extends, implements) - if chunk.class_metadata: - if chunk.class_metadata.get('extends'): - metadata['extends'] = chunk.class_metadata['extends'] - metadata['parent_types'] = chunk.class_metadata['extends'] - if chunk.class_metadata.get('implements'): - metadata['implements'] = chunk.class_metadata['implements'] - - # Try to extract additional metadata via regex patterns - patterns = METADATA_PATTERNS.get(chunk.language, {}) - - # Extract docstring - docstring = self._extract_docstring(chunk.content, chunk.language) - if docstring: - metadata['docstring'] = docstring[:500] - - # Extract signature - signature = self._extract_signature(chunk.content, chunk.language) - if signature: - metadata['signature'] = signature - - # Extract additional names not caught by AST - if not chunk.semantic_names: - names = [] - for pattern_type, pattern in patterns.items(): - matches = pattern.findall(chunk.content) - names.extend(matches) - if names: - metadata['semantic_names'] = list(set(names))[:10] - metadata['primary_name'] = names[0] - - # Fallback: Extract inheritance via regex if AST didn't find it - if 'extends' not in metadata and 'implements' not in metadata: - self._extract_inheritance_metadata(chunk.content, chunk.language, metadata) - - return metadata - - def _extract_inheritance_metadata( - self, - content: str, - language: str, - metadata: Dict[str, Any] - ) -> None: - """Extract inheritance, interfaces, and imports from code chunk""" - inheritance_patterns = CLASS_INHERITANCE_PATTERNS.get(language, {}) - - if not inheritance_patterns: - return - - # Extract extends (parent class) - if 'extends' in inheritance_patterns: - match = inheritance_patterns['extends'].search(content) - if match: - extends = match.group(1).strip() - # Clean up and split multiple classes (for multiple inheritance) - extends_list = [e.strip() for e in extends.split(',') if e.strip()] - if extends_list: - metadata['extends'] = extends_list - metadata['parent_types'] = extends_list # Alias for searchability - - # Extract implements (interfaces) - if 'implements' in inheritance_patterns: - match = inheritance_patterns['implements'].search(content) - if match: - implements = match.group(1).strip() - implements_list = [i.strip() for i in implements.split(',') if i.strip()] - if implements_list: - metadata['implements'] = implements_list - - # Extract imports/use statements - import_key = None - for key in ['import', 'use', 'using', 'require']: - if key in inheritance_patterns: - import_key = key - break - - if import_key: - matches = inheritance_patterns[import_key].findall(content) - if matches: - # Flatten if matches are tuples (from groups in regex) - imports = [] - for m in matches: - if isinstance(m, tuple): - imports.extend([x.strip() for x in m if x and x.strip()]) - else: - imports.append(m.strip()) - if imports: - metadata['imports'] = imports[:20] # Limit to 20 - - # Extract namespace/package - for key in ['namespace', 'package']: - if key in inheritance_patterns: - match = inheritance_patterns[key].search(content) - if match: - metadata[key] = match.group(1).strip() - break - - # Extract Rust impl for - if 'impl_for' in inheritance_patterns: - matches = inheritance_patterns['impl_for'].findall(content) - if matches: - # matches are tuples of (trait, type) - metadata['impl_traits'] = [m[0] for m in matches if m[0]] - metadata['impl_types'] = [m[1] for m in matches if m[1]] - - # Extract Scala with traits - if 'with' in inheritance_patterns: - matches = inheritance_patterns['with'].findall(content) - if matches: - metadata['with_traits'] = matches - - # Extract PHP type hints from docblocks - if 'type_hint' in inheritance_patterns: - matches = inheritance_patterns['type_hint'].findall(content) - if matches: - # Extract unique types referenced in docblocks - type_refs = list(set(matches))[:10] - metadata['type_references'] = type_refs - - # Extract PHTML template type hints (/** @var \Class $var */) - if 'template_type' in inheritance_patterns: - matches = inheritance_patterns['template_type'].findall(content) - if matches: - template_types = list(set(matches))[:10] - # Merge with type_references if exists - existing = metadata.get('type_references', []) - metadata['type_references'] = list(set(existing + template_types))[:15] - # Also add to related_classes for better searchability - metadata['related_classes'] = template_types - - # Extract PHP PHPDoc types (@param, @return, @throws) - if 'phpdoc_types' in inheritance_patterns: - matches = inheritance_patterns['phpdoc_types'].findall(content) - if matches: - # Filter and clean type names, handle union types - phpdoc_types = [] - for m in matches: - for t in m.split('|'): - t = t.strip().lstrip('\\') - if t and t[0].isupper(): # Only class names - phpdoc_types.append(t) - if phpdoc_types: - existing = metadata.get('type_references', []) - metadata['type_references'] = list(set(existing + phpdoc_types))[:20] - - def _extract_docstring(self, content: str, language: str) -> Optional[str]: - """Extract docstring from code chunk""" - if language == 'python': - match = re.search(r'"""([\s\S]*?)"""|\'\'\'([\s\S]*?)\'\'\'', content) - if match: - return (match.group(1) or match.group(2)).strip() - - elif language in ('javascript', 'typescript', 'java', 'kotlin', 'c_sharp', 'php', 'go', 'scala'): - match = re.search(r'/\*\*([\s\S]*?)\*/', content) - if match: - doc = match.group(1) - doc = re.sub(r'^\s*\*\s?', '', doc, flags=re.MULTILINE) - return doc.strip() - - elif language == 'rust': - lines = [] - for line in content.split('\n'): - if line.strip().startswith('///'): - lines.append(line.strip()[3:].strip()) - elif lines: - break - if lines: - return '\n'.join(lines) - - return None - - def _extract_signature(self, content: str, language: str) -> Optional[str]: - """Extract function/method signature from code chunk""" - lines = content.split('\n') - - for line in lines[:15]: - line = line.strip() - - if language == 'python': - if line.startswith(('def ', 'async def ', 'class ')): - sig = line - if line.startswith('class ') and ':' in line: - return line.split(':')[0] + ':' - if ')' not in sig and ':' not in sig: - idx = -1 - for i, l in enumerate(lines): - if l.strip() == line: - idx = i - break - if idx >= 0: - for next_line in lines[idx+1:idx+5]: - sig += ' ' + next_line.strip() - if ')' in next_line: - break - if ':' in sig: - return sig.split(':')[0] + ':' - return sig - - elif language in ('java', 'kotlin', 'c_sharp'): - if any(kw in line for kw in ['public ', 'private ', 'protected ', 'internal ', 'fun ']): - if '(' in line and not line.startswith('//'): - return line.split('{')[0].strip() - - elif language in ('javascript', 'typescript'): - if line.startswith(('function ', 'async function ', 'class ')): - return line.split('{')[0].strip() - if '=>' in line and '(' in line: - return line.split('=>')[0].strip() + ' =>' - - elif language == 'go': - if line.startswith('func ') or line.startswith('type '): - return line.split('{')[0].strip() - - elif language == 'rust': - if line.startswith(('fn ', 'pub fn ', 'async fn ', 'pub async fn ', 'impl ', 'struct ', 'trait ')): - return line.split('{')[0].strip() - - return None - - def _split_oversized_chunk( - self, - chunk: ASTChunk, - language: Optional[Language], - base_metadata: Dict[str, Any], - path: str - ) -> List[TextNode]: - """ - Split an oversized chunk using RecursiveCharacterTextSplitter. - - This is used when AST-parsed chunks (e.g., very large classes/functions) - still exceed the max_chunk_size. - """ - splitter = ( - self._get_text_splitter(language) - if language and language in AST_SUPPORTED_LANGUAGES - else self._default_splitter - ) - - sub_chunks = splitter.split_text(chunk.content) - nodes = [] - - # Parent ID for linking sub-chunks - parent_id = generate_deterministic_id(path, chunk.content, 0) - - for i, sub_chunk in enumerate(sub_chunks): - if not sub_chunk or not sub_chunk.strip(): - continue - - if len(sub_chunk.strip()) < self.min_chunk_size and len(sub_chunks) > 1: - continue - - metadata = dict(base_metadata) - metadata['content_type'] = ContentType.OVERSIZED_SPLIT.value - metadata['original_content_type'] = chunk.content_type.value - metadata['parent_chunk_id'] = parent_id - metadata['sub_chunk_index'] = i - metadata['total_sub_chunks'] = len(sub_chunks) - - # Preserve breadcrumb context - if chunk.parent_context: - metadata['parent_context'] = chunk.parent_context - metadata['parent_class'] = chunk.parent_context[-1] - - if chunk.semantic_names: - metadata['semantic_names'] = chunk.semantic_names - metadata['primary_name'] = chunk.semantic_names[0] - - # Deterministic ID for this sub-chunk - chunk_id = generate_deterministic_id(path, sub_chunk, i) - - node = TextNode( - id_=chunk_id, - text=sub_chunk, - metadata=metadata - ) - nodes.append(node) - - return nodes - - def split_documents(self, documents: List[LlamaDocument]) -> List[TextNode]: - """ - Split LlamaIndex documents using AST-based parsing. - - Args: - documents: List of LlamaIndex Document objects - - Returns: - List of TextNode objects with enriched metadata and deterministic IDs - """ - all_nodes = [] - - for doc in documents: - path = doc.metadata.get('path', 'unknown') - - # Determine Language enum - language = self._get_language_from_path(path) - - # Check if AST parsing is supported and beneficial - line_count = doc.text.count('\n') + 1 - use_ast = ( - language is not None - and language in AST_SUPPORTED_LANGUAGES - and line_count >= self.parser_threshold - and self._check_tree_sitter() - ) - - if use_ast: - nodes = self._split_with_ast(doc, language) - else: - nodes = self._split_fallback(doc, language) - - all_nodes.extend(nodes) - logger.debug(f"Split {path} into {len(nodes)} chunks (AST={use_ast})") - - return all_nodes - - def _split_with_ast( - self, - doc: LlamaDocument, - language: Language - ) -> List[TextNode]: - """Split document using AST parsing with breadcrumb context""" - text = doc.text - path = doc.metadata.get('path', 'unknown') - - # Try AST parsing - ast_chunks = self._parse_with_ast(text, language, path) - - if not ast_chunks: - return self._split_fallback(doc, language) - - nodes = [] - chunk_counter = 0 - - for ast_chunk in ast_chunks: - # Check if chunk is oversized - if len(ast_chunk.content) > self.max_chunk_size: - # Split oversized chunk - sub_nodes = self._split_oversized_chunk( - ast_chunk, - language, - doc.metadata, - path - ) - nodes.extend(sub_nodes) - chunk_counter += len(sub_nodes) - else: - # Create node with enriched metadata - metadata = self._extract_metadata(ast_chunk, doc.metadata) - metadata['chunk_index'] = chunk_counter - metadata['total_chunks'] = len(ast_chunks) - - # Deterministic ID - chunk_id = generate_deterministic_id(path, ast_chunk.content, chunk_counter) - - node = TextNode( - id_=chunk_id, - text=ast_chunk.content, - metadata=metadata - ) - nodes.append(node) - chunk_counter += 1 - - return nodes - - def _split_fallback( - self, - doc: LlamaDocument, - language: Optional[Language] = None - ) -> List[TextNode]: - """Fallback splitting using RecursiveCharacterTextSplitter""" - text = doc.text - path = doc.metadata.get('path', 'unknown') - - if not text or not text.strip(): - return [] - - splitter = ( - self._get_text_splitter(language) - if language and language in AST_SUPPORTED_LANGUAGES - else self._default_splitter - ) - - chunks = splitter.split_text(text) - nodes = [] - text_offset = 0 - - for i, chunk in enumerate(chunks): - if not chunk or not chunk.strip(): - continue - - if len(chunk.strip()) < self.min_chunk_size and len(chunks) > 1: - continue - - # Truncate if too large - if len(chunk) > 30000: - chunk = chunk[:30000] - - # Calculate line numbers - start_line = text[:text_offset].count('\n') + 1 if text_offset > 0 else 1 - chunk_pos = text.find(chunk, text_offset) - if chunk_pos >= 0: - text_offset = chunk_pos + len(chunk) - end_line = start_line + chunk.count('\n') - - # Extract metadata using regex patterns - lang_str = doc.metadata.get('language', 'text') - metadata = dict(doc.metadata) - metadata['content_type'] = ContentType.FALLBACK.value - metadata['chunk_index'] = i - metadata['total_chunks'] = len(chunks) - metadata['start_line'] = start_line - metadata['end_line'] = end_line - - # Try to extract semantic names - patterns = METADATA_PATTERNS.get(lang_str, {}) - names = [] - for pattern_type, pattern in patterns.items(): - matches = pattern.findall(chunk) - names.extend(matches) - if names: - metadata['semantic_names'] = list(set(names))[:10] - metadata['primary_name'] = names[0] - - # Extract class inheritance, interfaces, and imports (also for fallback) - self._extract_inheritance_metadata(chunk, lang_str, metadata) - - # Deterministic ID - chunk_id = generate_deterministic_id(path, chunk, i) - - node = TextNode( - id_=chunk_id, - text=chunk, - metadata=metadata - ) - nodes.append(node) - - return nodes - - @staticmethod - def get_supported_languages() -> List[str]: - """Return list of languages with AST support""" - return list(LANGUAGE_TO_TREESITTER.values()) - - @staticmethod - def is_ast_supported(path: str) -> bool: - """Check if AST parsing is supported for a file""" - ext = Path(path).suffix.lower() - lang = EXTENSION_TO_LANGUAGE.get(ext) - return lang is not None and lang in AST_SUPPORTED_LANGUAGES diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/chunking.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/chunking.py deleted file mode 100644 index 3c8cd8cc..00000000 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/chunking.py +++ /dev/null @@ -1,171 +0,0 @@ -from typing import List -import uuid -from llama_index.core.node_parser import SentenceSplitter -from llama_index.core.schema import Document, TextNode -from ..utils.utils import is_code_file - - -class CodeAwareSplitter: - """ - Code-aware text splitter that handles code and text differently. - - DEPRECATED: Use SemanticCodeSplitter instead, which provides: - - Full AST-aware parsing for multiple languages - - Better metadata extraction (docstrings, signatures, imports) - - Smarter chunk merging and boundary detection - - This class just wraps SentenceSplitter with different chunk sizes for - code vs text. For truly semantic code splitting, use SemanticCodeSplitter. - """ - - def __init__(self, code_chunk_size: int = 800, code_overlap: int = 200, - text_chunk_size: int = 1000, text_overlap: int = 200): - self.code_splitter = SentenceSplitter( - chunk_size=code_chunk_size, - chunk_overlap=code_overlap, - separator="\n\n", - ) - - self.text_splitter = SentenceSplitter( - chunk_size=text_chunk_size, - chunk_overlap=text_overlap, - ) - - def split_documents(self, documents: List[Document]) -> List[TextNode]: - """Split documents into chunks based on their language type""" - result = [] - - for doc in documents: - language = doc.metadata.get("language", "text") - is_code = is_code_file(language) - - splitter = self.code_splitter if is_code else self.text_splitter - - nodes = splitter.get_nodes_from_documents([doc]) - - for i, node in enumerate(nodes): - # Skip empty or whitespace-only chunks - if not node.text or not node.text.strip(): - continue - - # Truncate text if too large (>30k chars ≈ 7.5k tokens) - text = node.text - if len(text) > 30000: - text = text[:30000] - - metadata = dict(doc.metadata) - metadata["chunk_index"] = i - metadata["total_chunks"] = len(nodes) - - # Create TextNode with explicit UUID - chunk_node = TextNode( - id_=str(uuid.uuid4()), - text=text, - metadata=metadata - ) - result.append(chunk_node) - - return result - - def split_text_for_language(self, text: str, language: str) -> List[str]: - """Split text based on language type""" - is_code = is_code_file(language) - splitter = self.code_splitter if is_code else self.text_splitter - - temp_doc = Document(text=text, metadata={"language": language}) - nodes = splitter.get_nodes_from_documents([temp_doc]) - - return [node.text for node in nodes] - - -class FunctionAwareSplitter: - """ - Advanced splitter that tries to preserve function boundaries. - - DEPRECATED: Use SemanticCodeSplitter instead, which provides: - - Full AST-aware parsing for multiple languages - - Better metadata extraction (docstrings, signatures, imports) - - Smarter chunk merging and boundary detection - - This class is kept for backward compatibility only. - """ - - def __init__(self, max_chunk_size: int = 800, overlap: int = 200): - self.max_chunk_size = max_chunk_size - self.overlap = overlap - self.fallback_splitter = SentenceSplitter( - chunk_size=max_chunk_size, - chunk_overlap=overlap, - ) - - def split_by_functions(self, text: str, language: str) -> List[str]: - """Try to split code by functions/classes""" - - if language == 'python': - return self._split_python(text) - elif language in ['javascript', 'typescript', 'java', 'cpp', 'c', 'go', 'rust', 'php']: - return self._split_brace_language(text) - else: - temp_doc = Document(text=text) - nodes = self.fallback_splitter.get_nodes_from_documents([temp_doc]) - return [node.text for node in nodes] - - def _split_python(self, text: str) -> List[str]: - """Split Python code by top-level definitions""" - lines = text.split('\n') - chunks = [] - current_chunk = [] - - for line in lines: - stripped = line.lstrip() - - if stripped.startswith(('def ', 'class ', 'async def ')): - if current_chunk and len('\n'.join(current_chunk)) > 50: - chunks.append('\n'.join(current_chunk)) - current_chunk = [] - - current_chunk.append(line) - - if len('\n'.join(current_chunk)) > self.max_chunk_size: - chunks.append('\n'.join(current_chunk)) - current_chunk = [] - - if current_chunk: - chunks.append('\n'.join(current_chunk)) - - return chunks if chunks else [text] - - def _split_brace_language(self, text: str) -> List[str]: - """Split brace-based languages by functions/classes""" - chunks = [] - current_chunk = [] - brace_count = 0 - in_function = False - - lines = text.split('\n') - - for line in lines: - if any(keyword in line for keyword in - ['function ', 'class ', 'def ', 'fn ', 'func ', 'public ', 'private ', 'protected ']): - if '{' in line: - in_function = True - - current_chunk.append(line) - - brace_count += line.count('{') - line.count('}') - - if in_function and brace_count == 0 and len(current_chunk) > 3: - chunks.append('\n'.join(current_chunk)) - current_chunk = [] - in_function = False - - if len('\n'.join(current_chunk)) > self.max_chunk_size: - chunks.append('\n'.join(current_chunk)) - current_chunk = [] - brace_count = 0 - - if current_chunk: - chunks.append('\n'.join(current_chunk)) - - return chunks if chunks else [text] - diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/embedding_factory.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/embedding_factory.py new file mode 100644 index 00000000..899cd354 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/embedding_factory.py @@ -0,0 +1,93 @@ +""" +Embedding factory for creating embedding models based on configuration. +Supports switching between local (Ollama) and cloud (OpenRouter) providers. +""" + +import logging +from typing import Union + +from llama_index.core.base.embeddings.base import BaseEmbedding + +from ..models.config import RAGConfig +from .ollama_embedding import OllamaEmbedding +from .openrouter_embedding import OpenRouterEmbedding + +logger = logging.getLogger(__name__) + + +def create_embedding_model(config: RAGConfig) -> BaseEmbedding: + """ + Create an embedding model based on the configuration. + + Args: + config: RAGConfig with embedding provider settings + + Returns: + BaseEmbedding instance (OllamaEmbedding or OpenRouterEmbedding) + """ + provider = config.embedding_provider.lower() + + if provider == "ollama": + logger.info(f"Creating Ollama embedding model: {config.ollama_model}") + return OllamaEmbedding( + model=config.ollama_model, + base_url=config.ollama_base_url, + timeout=120.0, + expected_dim=config.embedding_dim + ) + + elif provider == "openrouter": + logger.info(f"Creating OpenRouter embedding model: {config.openrouter_model}") + return OpenRouterEmbedding( + api_key=config.openrouter_api_key, + model=config.openrouter_model, + api_base=config.openrouter_base_url, + timeout=60.0, + max_retries=3, + expected_dim=config.embedding_dim + ) + + else: + logger.warning(f"Unknown embedding provider '{provider}', defaulting to Ollama") + return OllamaEmbedding( + model=config.ollama_model, + base_url=config.ollama_base_url, + timeout=120.0, + expected_dim=config.embedding_dim + ) + + +def get_embedding_model_info(config: RAGConfig) -> dict: + """ + Get information about the configured embedding model. + + Args: + config: RAGConfig with embedding provider settings + + Returns: + Dictionary with provider info + """ + provider = config.embedding_provider.lower() + + if provider == "ollama": + return { + "provider": "ollama", + "model": config.ollama_model, + "base_url": config.ollama_base_url, + "embedding_dim": config.embedding_dim, + "type": "local" + } + elif provider == "openrouter": + return { + "provider": "openrouter", + "model": config.openrouter_model, + "base_url": config.openrouter_base_url, + "embedding_dim": config.embedding_dim, + "type": "cloud" + } + else: + return { + "provider": provider, + "embedding_dim": config.embedding_dim, + "type": "unknown" + } diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager.py deleted file mode 100644 index d5497d1a..00000000 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager.py +++ /dev/null @@ -1,1015 +0,0 @@ -from typing import Optional, List, Dict -from datetime import datetime, timezone -from pathlib import Path -import logging -import gc -import os -import time -import uuid -import hashlib - -from llama_index.core import VectorStoreIndex, StorageContext, Settings -from llama_index.core.schema import Document, TextNode -from llama_index.vector_stores.qdrant import QdrantVectorStore -from qdrant_client import QdrantClient -from qdrant_client.models import ( - Distance, VectorParams, Filter, FieldCondition, MatchAny, MatchValue, - CreateAlias, DeleteAlias, CreateAliasOperation, DeleteAliasOperation, - PointStruct -) - -from ..models.config import RAGConfig, IndexStats -from ..utils.utils import make_namespace, make_project_namespace -from .semantic_splitter import SemanticCodeSplitter -from .ast_splitter import ASTCodeSplitter -from .loader import DocumentLoader -from .openrouter_embedding import OpenRouterEmbedding - -logger = logging.getLogger(__name__) - -# Memory-efficient batch sizes -DOCUMENT_BATCH_SIZE = 50 # Process documents in batches to limit memory -INSERT_BATCH_SIZE = 50 # Batch size for LlamaIndex insert_nodes - - -class RAGIndexManager: - """Manage RAG indices for code repositories using Qdrant""" - - def __init__(self, config: RAGConfig): - self.config = config - - # Qdrant client for vector storage - self.qdrant_client = QdrantClient(url=config.qdrant_url) - logger.info(f"Connected to Qdrant at {config.qdrant_url}") - - # Configure OpenRouter embeddings - self.embed_model = OpenRouterEmbedding( - api_key=config.openrouter_api_key, - model=config.openrouter_model, - api_base=config.openrouter_base_url, - timeout=60.0, - max_retries=3, - expected_dim=config.embedding_dim - ) - - # Set global settings - Settings.embed_model = self.embed_model - Settings.chunk_size = config.chunk_size - Settings.chunk_overlap = config.chunk_overlap - - # Choose splitter based on environment variable or config - # AST splitter provides better semantic chunking for supported languages - use_ast_splitter = os.environ.get('RAG_USE_AST_SPLITTER', 'true').lower() == 'true' - - if use_ast_splitter: - logger.info("Using ASTCodeSplitter for code chunking (tree-sitter based)") - self.splitter = ASTCodeSplitter( - max_chunk_size=config.chunk_size, - min_chunk_size=min(200, config.chunk_size // 4), - chunk_overlap=config.chunk_overlap, - parser_threshold=10 # Minimum lines for AST parsing - ) - else: - logger.info("Using SemanticCodeSplitter for code chunking (regex-based)") - self.splitter = SemanticCodeSplitter( - max_chunk_size=config.chunk_size, - min_chunk_size=min(200, config.chunk_size // 4), - overlap=config.chunk_overlap - ) - - self.loader = DocumentLoader(config) - - def _get_project_collection_name(self, workspace: str, project: str) -> str: - """Generate Qdrant collection name from workspace/project (single collection per project)""" - namespace = make_project_namespace(workspace, project) - return f"{self.config.qdrant_collection_prefix}_{namespace}" - - def _get_collection_name(self, workspace: str, project: str, branch: str) -> str: - """Generate Qdrant collection name from workspace/project/branch (DEPRECATED - use _get_project_collection_name)""" - namespace = make_namespace(workspace, project, branch) - return f"{self.config.qdrant_collection_prefix}_{namespace}" - - def _generate_point_id(self, workspace: str, project: str, branch: str, path: str, chunk_index: int) -> str: - """Generate deterministic point ID for upsert (same content = same ID = replace)""" - key = f"{workspace}:{project}:{branch}:{path}:{chunk_index}" - return str(uuid.uuid5(uuid.NAMESPACE_DNS, key)) - - def _ensure_collection_exists(self, collection_name: str): - """Ensure Qdrant collection exists with proper configuration. - - If the collection_name is actually an alias, use the aliased collection instead. - """ - # First check if this is an alias - if self._alias_exists(collection_name): - logger.info(f"Collection name {collection_name} is an alias, using existing aliased collection") - return - - # Also check if there's a collection with this exact name - collections = self.qdrant_client.get_collections().collections - collection_names = [c.name for c in collections] - logger.debug(f"Existing collections: {collection_names}") - - if collection_name not in collection_names: - logger.info(f"Creating Qdrant collection: {collection_name}") - self.qdrant_client.create_collection( - collection_name=collection_name, - vectors_config=VectorParams( - size=self.config.embedding_dim, - distance=Distance.COSINE - ) - ) - logger.info(f"Created collection {collection_name}") - else: - logger.info(f"Collection {collection_name} already exists") - - def _get_storage_context(self, workspace: str, project: str, branch: str) -> StorageContext: - """Create storage context with Qdrant vector store (DEPRECATED - use direct Qdrant operations)""" - collection_name = self._get_project_collection_name(workspace, project) - self._ensure_collection_exists(collection_name) - - vector_store = QdrantVectorStore( - client=self.qdrant_client, - collection_name=collection_name, - enable_hybrid=False, # Disable hybrid mode to avoid string ID issues - batch_size=100 - ) - - return StorageContext.from_defaults(vector_store=vector_store) - - def _get_or_create_index( - self, - workspace: str, - project: str, - branch: str - ) -> VectorStoreIndex: - """Get or create vector index for the given namespace (DEPRECATED - use direct Qdrant operations)""" - namespace = make_namespace(workspace, project, branch) - logger.info(f"Getting/creating index for namespace: {namespace}") - - storage_context = self._get_storage_context(workspace, project, branch) - collection_name = self._get_project_collection_name(workspace, project) - - # Check if collection has data - collection_info = self.qdrant_client.get_collection(collection_name) - - if collection_info.points_count > 0: - logger.info(f"Loaded existing index for {namespace} ({collection_info.points_count} points)") - index = VectorStoreIndex.from_vector_store( - vector_store=storage_context.vector_store, - embed_model=self.embed_model - ) - else: - logger.info(f"Creating new index for {namespace}") - index = VectorStoreIndex.from_documents( - [], - storage_context=storage_context, - embed_model=self.embed_model, - show_progress=True - ) - - return index - - def _resolve_alias_to_collection(self, alias_name: str) -> Optional[str]: - """Resolve an alias to its underlying collection name""" - try: - # Get all aliases and find the one matching our alias_name - aliases = self.qdrant_client.get_aliases() - for alias in aliases.aliases: - if alias.alias_name == alias_name: - return alias.collection_name - except Exception as e: - logger.debug(f"Error resolving alias {alias_name}: {e}") - return None - - def _alias_exists(self, alias_name: str) -> bool: - """Check if an alias exists""" - try: - aliases = self.qdrant_client.get_aliases() - exists = any(a.alias_name == alias_name for a in aliases.aliases) - logger.debug(f"Checking if alias '{alias_name}' exists: {exists}. All aliases: {[a.alias_name for a in aliases.aliases]}") - return exists - except Exception as e: - logger.warning(f"Error checking alias {alias_name}: {e}") - return False - - def estimate_repository_size( - self, - repo_path: str, - exclude_patterns: Optional[List[str]] = None - ) -> tuple[int, int]: - """Estimate repository size (file count and chunk count) without actually indexing. - - Uses streaming approach to avoid loading all files into memory. - - Args: - repo_path: Path to the repository - exclude_patterns: Additional patterns to exclude - - Returns: - Tuple of (file_count, estimated_chunk_count) - """ - logger.info(f"Estimating repository size for: {repo_path}") - - repo_path_obj = Path(repo_path) - - # Count files without loading content (memory-efficient) - file_list = list(self.loader.iter_repository_files(repo_path_obj, exclude_patterns)) - file_count = len(file_list) - logger.info(f"Found {file_count} files for estimation") - - if file_count == 0: - return 0, 0 - - # Sample-based chunk estimation for large repos - # For repos with many files, sample a subset and extrapolate - SAMPLE_SIZE = 100 - chunk_count = 0 - - if file_count <= SAMPLE_SIZE: - # Small repo: count all chunks - for i in range(0, file_count, DOCUMENT_BATCH_SIZE): - batch = file_list[i:i + DOCUMENT_BATCH_SIZE] - documents = self.loader.load_file_batch( - batch, repo_path_obj, "estimate", "estimate", "estimate", "estimate" - ) - if documents: - chunks = self.splitter.split_documents(documents) - chunk_count += len(chunks) - del chunks - del documents - gc.collect() - else: - # Large repo: sample and extrapolate - import random - sample_files = random.sample(file_list, SAMPLE_SIZE) - sample_chunk_count = 0 - - for i in range(0, len(sample_files), DOCUMENT_BATCH_SIZE): - batch = sample_files[i:i + DOCUMENT_BATCH_SIZE] - documents = self.loader.load_file_batch( - batch, repo_path_obj, "estimate", "estimate", "estimate", "estimate" - ) - if documents: - chunks = self.splitter.split_documents(documents) - sample_chunk_count += len(chunks) - del chunks - del documents - - # Extrapolate - avg_chunks_per_file = sample_chunk_count / SAMPLE_SIZE - chunk_count = int(avg_chunks_per_file * file_count) - logger.info(f"Estimated ~{avg_chunks_per_file:.1f} chunks/file from {SAMPLE_SIZE} samples") - gc.collect() - - logger.info(f"Estimated {chunk_count} chunks from {file_count} files") - - return file_count, chunk_count - - def index_repository( - self, - repo_path: str, - workspace: str, - project: str, - branch: str, - commit: str, - exclude_patterns: Optional[List[str]] = None - ) -> IndexStats: - """Index entire repository for a branch using single-collection-per-project architecture. - - Uses versioned collections with alias-based atomic swap for zero-downtime reindexing: - 1. Create a new versioned collection (e.g., project_v1234567) - 2. Index all files into the temp collection - 3. On success, atomically swap the alias to point to new collection - 4. Delete the old versioned collection - - This ensures the old index remains available if indexing fails. - Branch metadata is stored in point payloads for multi-branch queries. - """ - logger.info(f"Indexing repository: {workspace}/{project}/{branch} from {repo_path}") - - # Convert string to Path object - repo_path_obj = Path(repo_path) - - # Project-level alias name (no branch in name) - alias_name = self._get_project_collection_name(workspace, project) - temp_collection_name = f"{alias_name}_v{int(time.time())}" - - # Check if alias or collection exists and get existing branch data to preserve - old_collection_exists = self._alias_exists(alias_name) - existing_other_branch_points = [] - - if not old_collection_exists: - collections = self.qdrant_client.get_collections().collections - collection_names = [c.name for c in collections] - old_collection_exists = alias_name in collection_names - - # If collection exists, preserve points from OTHER branches (not the one being reindexed) - if old_collection_exists: - try: - actual_collection = self._resolve_alias_to_collection(alias_name) or alias_name - logger.info(f"Preserving points from other branches in {actual_collection}...") - - # Scroll through all points NOT in the branch being indexed - offset = None - while True: - results = self.qdrant_client.scroll( - collection_name=actual_collection, - limit=100, - offset=offset, - scroll_filter=Filter( - must_not=[ - FieldCondition( - key="branch", - match=MatchValue(value=branch) - ) - ] - ), - with_payload=True, - with_vectors=True - ) - points, next_offset = results - existing_other_branch_points.extend(points) - - if next_offset is None or len(points) < 100: - break - offset = next_offset - - logger.info(f"Found {len(existing_other_branch_points)} points from other branches to preserve") - except Exception as e: - logger.warning(f"Could not read existing points: {e}") - - # Clean up any leftover versioned collections from previous failed attempts - collections = self.qdrant_client.get_collections().collections - collection_names = [c.name for c in collections] - for coll_name in collection_names: - if coll_name.startswith(f"{alias_name}_v") and coll_name != temp_collection_name: - current_alias_target = self._resolve_alias_to_collection(alias_name) - if current_alias_target != coll_name: - logger.info(f"Cleaning up orphaned versioned collection: {coll_name}") - try: - self.qdrant_client.delete_collection(coll_name) - except Exception as e: - logger.warning(f"Failed to clean up orphaned collection: {e}") - - # Create new temporary collection - logger.info(f"Creating temporary collection: {temp_collection_name}") - self.qdrant_client.create_collection( - collection_name=temp_collection_name, - vectors_config=VectorParams( - size=self.config.embedding_dim, - distance=Distance.COSINE - ) - ) - - # Get file list using DocumentLoader (low memory) - file_list = list(self.loader.iter_repository_files(repo_path_obj, exclude_patterns)) - total_files = len(file_list) - logger.info(f"Found {total_files} files to index for branch '{branch}'") - - if total_files == 0: - logger.warning("No documents to index") - self.qdrant_client.delete_collection(temp_collection_name) - return self._get_branch_index_stats(workspace, project, branch) - - # Validate file limit first (BEFORE any expensive operations) - if self.config.max_files_per_index > 0 and total_files > self.config.max_files_per_index: - self.qdrant_client.delete_collection(temp_collection_name) - raise ValueError( - f"Repository exceeds file limit: {total_files} files (max: {self.config.max_files_per_index}). " - f"Use exclude patterns in Project Settings → RAG Indexing to exclude unnecessary directories." - ) - - # Estimate chunk count and validate BEFORE starting embeddings - # This prevents wasting API calls on repos that will fail the chunk limit - if self.config.max_chunks_per_index > 0: - logger.info("Estimating chunk count before indexing...") - _, estimated_chunks = self.estimate_repository_size(repo_path, exclude_patterns) - logger.info(f"Estimated chunks: {estimated_chunks}, limit: {self.config.max_chunks_per_index}") - - # Add 20% buffer for estimation variance - if estimated_chunks > self.config.max_chunks_per_index * 1.2: - self.qdrant_client.delete_collection(temp_collection_name) - raise ValueError( - f"Repository estimated to exceed chunk limit: ~{estimated_chunks} chunks (max: {self.config.max_chunks_per_index}). " - f"Use exclude patterns in Project Settings → RAG Indexing to exclude large directories." - ) - - document_count = 0 - chunk_count = 0 - successful_chunks = 0 - failed_chunks = 0 - - try: - # First, copy preserved points from other branches to temp collection - if existing_other_branch_points: - logger.info(f"Copying {len(existing_other_branch_points)} points from other branches...") - for i in range(0, len(existing_other_branch_points), INSERT_BATCH_SIZE): - batch = existing_other_branch_points[i:i + INSERT_BATCH_SIZE] - points_to_upsert = [ - PointStruct( - id=p.id, - vector=p.vector, - payload=p.payload - ) for p in batch - ] - self.qdrant_client.upsert( - collection_name=temp_collection_name, - points=points_to_upsert - ) - logger.info("Other branch points copied successfully") - - # MEMORY-EFFICIENT STREAMING: Process documents in batches - logger.info("Starting memory-efficient streaming indexing...") - - # Process files in batches - batch_num = 0 - total_batches = (total_files + DOCUMENT_BATCH_SIZE - 1) // DOCUMENT_BATCH_SIZE - - for i in range(0, total_files, DOCUMENT_BATCH_SIZE): - batch_num += 1 - file_batch = file_list[i:i + DOCUMENT_BATCH_SIZE] - - # Load batch of documents using DocumentLoader - documents = self.loader.load_file_batch( - file_batch, repo_path_obj, workspace, project, branch, commit - ) - document_count += len(documents) - - if not documents: - continue - - # Split documents into chunks - chunks = self.splitter.split_documents(documents) - batch_chunk_count = len(chunks) - chunk_count += batch_chunk_count - - # Check chunk limit - if self.config.max_chunks_per_index > 0 and chunk_count > self.config.max_chunks_per_index: - self.qdrant_client.delete_collection(temp_collection_name) - raise ValueError( - f"Repository exceeds chunk limit: {chunk_count}+ chunks (max: {self.config.max_chunks_per_index}). " - f"Use exclude patterns to exclude large directories." - ) - - # Prepare points with deterministic IDs for upsert - # First, collect all chunks with their metadata - chunk_data = [] # List of (point_id, chunk, metadata) - chunks_by_file: Dict[str, List[TextNode]] = {} - for chunk in chunks: - path = chunk.metadata.get("path", "unknown") - if path not in chunks_by_file: - chunks_by_file[path] = [] - chunks_by_file[path].append(chunk) - - for path, file_chunks in chunks_by_file.items(): - for chunk_index, chunk in enumerate(file_chunks): - point_id = self._generate_point_id(workspace, project, branch, path, chunk_index) - chunk.metadata["indexed_at"] = datetime.now(timezone.utc).isoformat() - chunk_data.append((point_id, chunk)) - - # Batch embed all chunks at once (much more efficient) - texts_to_embed = [chunk.text for _, chunk in chunk_data] - embeddings = self.embed_model.get_text_embedding_batch(texts_to_embed) - - # Build points with embeddings - points = [] - for (point_id, chunk), embedding in zip(chunk_data, embeddings): - points.append(PointStruct( - id=point_id, - vector=embedding, - payload={ - **chunk.metadata, - "text": chunk.text, - "_node_content": chunk.text, - } - )) - - # Upsert in batches (idempotent - same ID = replace) - try: - for j in range(0, len(points), INSERT_BATCH_SIZE): - insert_batch = points[j:j + INSERT_BATCH_SIZE] - self.qdrant_client.upsert( - collection_name=temp_collection_name, - points=insert_batch - ) - successful_chunks += batch_chunk_count - except Exception as e: - logger.error(f"Failed to upsert batch {batch_num}: {e}") - failed_chunks += batch_chunk_count - - logger.info( - f"Batch {batch_num}/{total_batches}: processed {len(documents)} files, " - f"{batch_chunk_count} chunks" - ) - - # CRITICAL: Free memory after each batch - del documents - del chunks - del points - - # Aggressive garbage collection every few batches - if batch_num % 5 == 0: - gc.collect() - logger.debug(f"Memory cleanup after batch {batch_num}") - - logger.info( - f"Streaming indexing complete: {document_count} files, " - f"{successful_chunks}/{chunk_count} chunks indexed ({failed_chunks} failed)" - ) - - # Verify temp collection has data - temp_collection_info = self.qdrant_client.get_collection(temp_collection_name) - if temp_collection_info.points_count == 0: - raise Exception("Temporary collection is empty after indexing") - - # SUCCESS: Atomic alias swap (zero-downtime) - logger.info(f"Indexing successful. Performing atomic alias swap...") - - # Check if there's a collection (not alias) with the target name - collections = self.qdrant_client.get_collections().collections - collection_names = [c.name for c in collections] - is_direct_collection = alias_name in collection_names and not self._alias_exists(alias_name) - - # Find old versioned collection BEFORE alias operations - old_versioned_name = None - if old_collection_exists and not is_direct_collection: - old_versioned_name = self._resolve_alias_to_collection(alias_name) - - alias_operations = [] - - # Delete old alias if exists (not direct collection) - if old_collection_exists and not is_direct_collection: - alias_operations.append( - DeleteAliasOperation(delete_alias=DeleteAlias(alias_name=alias_name)) - ) - - # Create new alias pointing to temp collection - alias_operations.append( - CreateAliasOperation(create_alias=CreateAlias( - alias_name=alias_name, - collection_name=temp_collection_name - )) - ) - - # Perform atomic alias swap - self.qdrant_client.update_collection_aliases( - change_aliases_operations=alias_operations - ) - - logger.info(f"Alias swap completed successfully: {alias_name} -> {temp_collection_name}") - - # NOW delete old collections (after alias swap is complete) - if is_direct_collection: - logger.info(f"Migrating from direct collection to alias-based indexing. Deleting old collection: {alias_name}") - try: - self.qdrant_client.delete_collection(alias_name) - except Exception as del_err: - logger.warning(f"Failed to delete old direct collection: {del_err}") - elif old_versioned_name and old_versioned_name != temp_collection_name: - logger.info(f"Deleting old versioned collection: {old_versioned_name}") - try: - self.qdrant_client.delete_collection(old_versioned_name) - except Exception as del_err: - logger.warning(f"Failed to delete old collection: {del_err}") - - except Exception as e: - logger.error(f"Indexing failed: {e}") - logger.info(f"Cleaning up temporary collection, old index preserved") - try: - self.qdrant_client.delete_collection(temp_collection_name) - except Exception as cleanup_error: - logger.warning(f"Failed to clean up temp collection: {cleanup_error}") - raise e - finally: - # Free the preserved points from memory - del existing_other_branch_points - # Force garbage collection - gc.collect() - logger.info("Memory cleanup completed after indexing") - - # Store metadata - self._store_metadata(workspace, project, branch, commit, [], [], document_count, chunk_count) - - # Return stats with actual counts from this indexing run - namespace = make_namespace(workspace, project, branch) - return IndexStats( - namespace=namespace, - document_count=document_count, - chunk_count=successful_chunks, - last_updated=datetime.now(timezone.utc).isoformat(), - workspace=workspace, - project=project, - branch=branch - ) - - def _store_metadata( - self, - workspace: str, - project: str, - branch: str, - commit: str, - documents: List[Document], - chunks: List[TextNode], - document_count: Optional[int] = None, - chunk_count: Optional[int] = None - ): - """Store metadata in Qdrant collection payload""" - namespace = make_namespace(workspace, project, branch) - collection_name = self._get_project_collection_name(workspace, project) - - # Use provided counts or calculate from lists - doc_count = document_count if document_count is not None else len(documents) - chk_count = chunk_count if chunk_count is not None else len(chunks) - - # Update collection metadata - metadata = { - "namespace": namespace, - "workspace": workspace, - "project": project, - "branch": branch, - "commit": commit, - "document_count": doc_count, - "chunk_count": chk_count, - "last_updated": datetime.now(timezone.utc).isoformat(), - } - - # Store in collection's payload_schema or as a special point - # For now, we'll just log it - Qdrant will track points automatically - logger.info(f"Indexed {namespace}: {doc_count} docs, {chk_count} chunks") - - def update_files( - self, - file_paths: List[str], - repo_base: str, - workspace: str, - project: str, - branch: str, - commit: str - ) -> IndexStats: - """Update specific files in the index for a specific branch (Delete Old -> Insert New). - - Uses single project collection with branch in metadata. - Deterministic point IDs ensure same file+branch+chunk = replace, not duplicate. - """ - logger.info(f"Updating {len(file_paths)} files in {workspace}/{project} for branch '{branch}'") - - # 1. PREPARATION - repo_base_obj = Path(repo_base) - file_path_objs = [Path(fp) for fp in file_paths] - collection_name = self._get_project_collection_name(workspace, project) - - # Ensure collection exists - self._ensure_collection_exists(collection_name) - - # 2. DELETE OLD CHUNKS for these files AND this branch - # Only delete points for the specific branch being updated - logger.info(f"Purging existing vectors for {len(file_paths)} files in branch '{branch}'...") - - try: - self.qdrant_client.delete( - collection_name=collection_name, - points_selector=Filter( - must=[ - FieldCondition( - key="path", - match=MatchAny(any=file_paths) - ), - FieldCondition( - key="branch", - match=MatchValue(value=branch) - ) - ] - ) - ) - except Exception as e: - logger.error(f"Error deleting old chunks: {e}") - raise e - - # 3. LOAD & SPLIT NEW CONTENT - documents = self.loader.load_specific_files( - file_paths=file_path_objs, - repo_base=repo_base_obj, - workspace=workspace, - project=project, - branch=branch, - commit=commit - ) - - if not documents: - logger.warning("No documents loaded from provided paths.") - return self._get_project_index_stats(workspace, project) - - chunks = self.splitter.split_documents(documents) - logger.info(f"Generated {len(chunks)} new chunks") - - # 4. INSERT NEW CHUNKS with deterministic IDs - # Group chunks by file path to assign chunk indices - chunks_by_file: Dict[str, List[TextNode]] = {} - for chunk in chunks: - path = chunk.metadata.get("path", "unknown") - if path not in chunks_by_file: - chunks_by_file[path] = [] - chunks_by_file[path].append(chunk) - - # Collect all chunks with their metadata - chunk_data = [] # List of (point_id, chunk) - for path, file_chunks in chunks_by_file.items(): - for chunk_index, chunk in enumerate(file_chunks): - point_id = self._generate_point_id(workspace, project, branch, path, chunk_index) - chunk.metadata["indexed_at"] = datetime.now(timezone.utc).isoformat() - chunk_data.append((point_id, chunk)) - - # Batch embed all chunks at once (much more efficient) - texts_to_embed = [chunk.text for _, chunk in chunk_data] - embeddings = self.embed_model.get_text_embedding_batch(texts_to_embed) - - # Build points with embeddings - points = [] - for (point_id, chunk), embedding in zip(chunk_data, embeddings): - points.append(PointStruct( - id=point_id, - vector=embedding, - payload={ - **chunk.metadata, - "text": chunk.text, - "_node_content": chunk.text, - } - )) - - # Upsert in batches - for i in range(0, len(points), INSERT_BATCH_SIZE): - batch = points[i:i + INSERT_BATCH_SIZE] - try: - self.qdrant_client.upsert( - collection_name=collection_name, - points=batch - ) - except Exception as e: - logger.error(f"Failed to upsert batch {i // INSERT_BATCH_SIZE}: {e}") - raise e - - logger.info(f"Successfully updated {len(chunks)} chunks for branch '{branch}'") - - return self._get_project_index_stats(workspace, project) - - def delete_files(self, file_paths: List[str], workspace: str, project: str, branch: str) -> IndexStats: - """Delete specific files from the index for a specific branch""" - logger.info(f"Deleting {len(file_paths)} files from {workspace}/{project} branch '{branch}'") - collection_name = self._get_project_collection_name(workspace, project) - - try: - self.qdrant_client.delete( - collection_name=collection_name, - points_selector=Filter( - must=[ - FieldCondition( - key="path", - match=MatchAny(any=file_paths) - ), - FieldCondition( - key="branch", - match=MatchValue(value=branch) - ) - ] - ) - ) - logger.info(f"Deleted {len(file_paths)} files from branch '{branch}'") - except Exception as e: - logger.warning(f"Error deleting files: {e}") - - return self._get_project_index_stats(workspace, project) - - def delete_branch(self, workspace: str, project: str, branch: str) -> bool: - """Delete all points for a specific branch from the project collection. - - Used when a branch is deleted or cleaned up. - Does NOT delete the entire collection - only points for this branch. - """ - logger.info(f"Deleting all points for branch '{branch}' from {workspace}/{project}") - collection_name = self._get_project_collection_name(workspace, project) - - try: - # Check if collection exists - collections = [c.name for c in self.qdrant_client.get_collections().collections] - if collection_name not in collections and not self._alias_exists(collection_name): - logger.warning(f"Collection {collection_name} does not exist") - return False - - self.qdrant_client.delete( - collection_name=collection_name, - points_selector=Filter( - must=[ - FieldCondition( - key="branch", - match=MatchValue(value=branch) - ) - ] - ) - ) - logger.info(f"Successfully deleted all points for branch '{branch}'") - return True - except Exception as e: - logger.error(f"Failed to delete branch '{branch}': {e}") - return False - - def get_branch_point_count(self, workspace: str, project: str, branch: str) -> int: - """Get the number of points for a specific branch.""" - collection_name = self._get_project_collection_name(workspace, project) - - try: - # Check if collection exists - collections = [c.name for c in self.qdrant_client.get_collections().collections] - if collection_name not in collections and not self._alias_exists(collection_name): - return 0 - - result = self.qdrant_client.count( - collection_name=collection_name, - count_filter=Filter( - must=[ - FieldCondition( - key="branch", - match=MatchValue(value=branch) - ) - ] - ) - ) - return result.count - except Exception as e: - logger.error(f"Failed to get point count for branch '{branch}': {e}") - return 0 - - def get_indexed_branches(self, workspace: str, project: str) -> List[str]: - """Get list of branches that have points in the collection.""" - collection_name = self._get_project_collection_name(workspace, project) - - try: - # Check if collection exists - collections = [c.name for c in self.qdrant_client.get_collections().collections] - if collection_name not in collections and not self._alias_exists(collection_name): - return [] - - # Scroll through points and collect unique branches - # This is a simplified approach - for large collections, consider using facets - branches = set() - offset = None - limit = 100 - - while True: - results = self.qdrant_client.scroll( - collection_name=collection_name, - limit=limit, - offset=offset, - with_payload=["branch"], - with_vectors=False - ) - - points, next_offset = results - - for point in points: - if point.payload and "branch" in point.payload: - branches.add(point.payload["branch"]) - - if next_offset is None or len(points) < limit: - break - offset = next_offset - - return list(branches) - except Exception as e: - logger.error(f"Failed to get indexed branches: {e}") - return [] - - def _get_project_index_stats(self, workspace: str, project: str) -> IndexStats: - """Get statistics about a project's index (all branches combined)""" - collection_name = self._get_project_collection_name(workspace, project) - namespace = make_project_namespace(workspace, project) - - try: - collection_info = self.qdrant_client.get_collection(collection_name) - chunk_count = collection_info.points_count - - return IndexStats( - namespace=namespace, - document_count=0, - chunk_count=chunk_count, - last_updated=datetime.now(timezone.utc).isoformat(), - workspace=workspace, - project=project, - branch="*" # Indicates all branches - ) - except Exception: - return IndexStats( - namespace=namespace, - document_count=0, - chunk_count=0, - last_updated="", - workspace=workspace, - project=project, - branch="*" - ) - - def delete_index(self, workspace: str, project: str, branch: str): - """Delete branch data from project index. - - If branch is specified, only deletes that branch's points. - To delete entire project collection, use delete_project_index(). - """ - if branch and branch != "*": - # Delete only this branch's points - self.delete_branch(workspace, project, branch) - else: - # Delete all branches - kept for backward compatibility - self.delete_project_index(workspace, project) - - def delete_project_index(self, workspace: str, project: str): - """Delete entire project collection (all branches)""" - collection_name = self._get_project_collection_name(workspace, project) - namespace = make_project_namespace(workspace, project) - - logger.info(f"Deleting entire project index for {namespace}") - - # Delete Qdrant collection and any aliases - try: - # Check if it's an alias - if self._alias_exists(collection_name): - actual_collection = self._resolve_alias_to_collection(collection_name) - # Delete alias first - self.qdrant_client.delete_alias(collection_name) - # Then delete actual collection - if actual_collection: - self.qdrant_client.delete_collection(actual_collection) - else: - self.qdrant_client.delete_collection(collection_name) - logger.info(f"Deleted Qdrant collection: {collection_name}") - except Exception as e: - logger.warning(f"Failed to delete Qdrant collection: {e}") - - def _get_index_stats(self, workspace: str, project: str, branch: str) -> IndexStats: - """Get statistics about a branch index (for backward compatibility)""" - return self._get_branch_index_stats(workspace, project, branch) - - def _get_branch_index_stats(self, workspace: str, project: str, branch: str) -> IndexStats: - """Get statistics about a specific branch within a project collection""" - namespace = make_namespace(workspace, project, branch) - collection_name = self._get_project_collection_name(workspace, project) - - try: - # Count points for this specific branch - count_result = self.qdrant_client.count( - collection_name=collection_name, - count_filter=Filter( - must=[ - FieldCondition( - key="branch", - match=MatchValue(value=branch) - ) - ] - ) - ) - chunk_count = count_result.count - - return IndexStats( - namespace=namespace, - document_count=0, - chunk_count=chunk_count, - last_updated=datetime.now(timezone.utc).isoformat(), - workspace=workspace, - project=project, - branch=branch - ) - except Exception: - return IndexStats( - namespace=namespace, - document_count=0, - chunk_count=0, - last_updated="", - workspace=workspace, - project=project, - branch=branch - ) - - def list_indices(self) -> List[IndexStats]: - """List all project indices with branch breakdown""" - indices = [] - - # Get all collections from Qdrant - collections = self.qdrant_client.get_collections().collections - - for collection in collections: - # Parse collection name to extract workspace/project (new format: no branch) - if collection.name.startswith(f"{self.config.qdrant_collection_prefix}_"): - namespace = collection.name[len(f"{self.config.qdrant_collection_prefix}_"):] - parts = namespace.split("__") - - if len(parts) == 2: - # New format: workspace__project - workspace, project = parts - stats = self._get_project_index_stats(workspace, project) - indices.append(stats) - elif len(parts) == 3: - # Legacy format: workspace__project__branch (for migration) - workspace, project, branch = parts - stats = self._get_branch_index_stats(workspace, project, branch) - indices.append(stats) - - return indices - diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/__init__.py new file mode 100644 index 00000000..d06d664f --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/__init__.py @@ -0,0 +1,10 @@ +""" +Index Manager module for RAG indexing operations. + +Provides RAGIndexManager as the main entry point for indexing repositories, +managing collections, and handling branch-level operations. +""" + +from .manager import RAGIndexManager + +__all__ = ["RAGIndexManager"] diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/branch_manager.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/branch_manager.py new file mode 100644 index 00000000..5ea50f45 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/branch_manager.py @@ -0,0 +1,242 @@ +""" +Branch-level operations for RAG indices. + +Handles branch-specific point management within project collections. +""" + +import logging +from typing import List, Set, Optional, Iterator, Generator + +from qdrant_client import QdrantClient +from qdrant_client.models import Filter, FieldCondition, MatchValue, PointStruct + +logger = logging.getLogger(__name__) + + +class BranchManager: + """Manages branch-level operations within project collections.""" + + def __init__(self, client: QdrantClient): + self.client = client + + def delete_branch_points( + self, + collection_name: str, + branch: str + ) -> bool: + """Delete all points for a specific branch from the collection.""" + logger.info(f"Deleting all points for branch '{branch}' from {collection_name}") + + try: + self.client.delete( + collection_name=collection_name, + points_selector=Filter( + must=[ + FieldCondition( + key="branch", + match=MatchValue(value=branch) + ) + ] + ) + ) + logger.info(f"Successfully deleted all points for branch '{branch}'") + return True + except Exception as e: + logger.error(f"Failed to delete branch '{branch}': {e}") + return False + + def get_branch_point_count( + self, + collection_name: str, + branch: str + ) -> int: + """Get the number of points for a specific branch.""" + try: + result = self.client.count( + collection_name=collection_name, + count_filter=Filter( + must=[ + FieldCondition( + key="branch", + match=MatchValue(value=branch) + ) + ] + ) + ) + return result.count + except Exception as e: + logger.error(f"Failed to get point count for branch '{branch}': {e}") + return 0 + + def get_indexed_branches(self, collection_name: str) -> List[str]: + """Get list of branches that have points in the collection.""" + try: + branches: Set[str] = set() + offset = None + limit = 100 + + while True: + results = self.client.scroll( + collection_name=collection_name, + limit=limit, + offset=offset, + with_payload=["branch"], + with_vectors=False + ) + + points, next_offset = results + + for point in points: + if point.payload and "branch" in point.payload: + branches.add(point.payload["branch"]) + + if next_offset is None or len(points) < limit: + break + offset = next_offset + + return list(branches) + except Exception as e: + logger.error(f"Failed to get indexed branches: {e}") + return [] + + def preserve_other_branch_points( + self, + collection_name: str, + exclude_branch: str, + batch_size: int = 100 + ) -> Generator[List[PointStruct], None, None]: + """Stream points from branches other than the one being reindexed. + + Yields batches of points instead of loading all into memory. + Used during full reindex to keep data from other branches. + """ + logger.info(f"Streaming points from branches other than '{exclude_branch}'...") + + offset = None + total_yielded = 0 + + try: + while True: + results = self.client.scroll( + collection_name=collection_name, + limit=batch_size, + offset=offset, + scroll_filter=Filter( + must_not=[ + FieldCondition( + key="branch", + match=MatchValue(value=exclude_branch) + ) + ] + ), + with_payload=True, + with_vectors=True + ) + points, next_offset = results + + if points: + total_yielded += len(points) + yield points + + if next_offset is None or len(points) < batch_size: + break + offset = next_offset + + logger.info(f"Streamed {total_yielded} points from other branches") + except Exception as e: + logger.warning(f"Could not read existing points: {e}") + return + + def stream_copy_points_to_collection( + self, + source_collection: str, + target_collection: str, + exclude_branch: str, + batch_size: int = 50 + ) -> int: + """Stream copy points from one collection to another, excluding a branch. + + Memory-efficient alternative to preserve_other_branch_points + copy_points_to_collection. + Skips copying if vector dimensions don't match between collections. + """ + # Check vector dimensions match before copying + try: + source_info = self.client.get_collection(source_collection) + target_info = self.client.get_collection(target_collection) + + # Get dimensions from vector config + source_dim = None + target_dim = None + + if hasattr(source_info.config.params, 'vectors'): + vectors_config = source_info.config.params.vectors + if hasattr(vectors_config, 'size'): + source_dim = vectors_config.size + elif isinstance(vectors_config, dict) and '' in vectors_config: + source_dim = vectors_config[''].size + + if hasattr(target_info.config.params, 'vectors'): + vectors_config = target_info.config.params.vectors + if hasattr(vectors_config, 'size'): + target_dim = vectors_config.size + elif isinstance(vectors_config, dict) and '' in vectors_config: + target_dim = vectors_config[''].size + + if source_dim and target_dim and source_dim != target_dim: + logger.warning( + f"Skipping branch preservation: dimension mismatch " + f"(source: {source_dim}, target: {target_dim}). " + f"Re-embedding required for all branches." + ) + return 0 + + except Exception as e: + logger.warning(f"Could not verify collection dimensions: {e}") + # Continue anyway - will fail at upsert if dimensions don't match + total_copied = 0 + + for batch in self.preserve_other_branch_points(source_collection, exclude_branch, batch_size): + points_to_upsert = [ + PointStruct( + id=p.id, + vector=p.vector, + payload=p.payload + ) for p in batch + ] + self.client.upsert( + collection_name=target_collection, + points=points_to_upsert + ) + total_copied += len(points_to_upsert) + + if total_copied > 0: + logger.info(f"Copied {total_copied} points to {target_collection}") + return total_copied + + def copy_points_to_collection( + self, + points: List, + target_collection: str, + batch_size: int = 50 + ) -> None: + """Copy preserved points to a new collection.""" + if not points: + return + + logger.info(f"Copying {len(points)} points to {target_collection}...") + + for i in range(0, len(points), batch_size): + batch = points[i:i + batch_size] + points_to_upsert = [ + PointStruct( + id=p.id, + vector=p.vector, + payload=p.payload + ) for p in batch + ] + self.client.upsert( + collection_name=target_collection, + points=points_to_upsert + ) + + logger.info("Points copied successfully") diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/collection_manager.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/collection_manager.py new file mode 100644 index 00000000..039f06bf --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/collection_manager.py @@ -0,0 +1,187 @@ +""" +Qdrant collection and alias management utilities. + +Handles collection creation, alias operations, and resolution. +""" + +import logging +import time +from typing import Optional, List + +from qdrant_client import QdrantClient +from qdrant_client.models import ( + Distance, VectorParams, + CreateAlias, DeleteAlias, CreateAliasOperation, DeleteAliasOperation, + PayloadSchemaType, TextIndexParams, TokenizerType +) + +logger = logging.getLogger(__name__) + + +class CollectionManager: + """Manages Qdrant collections and aliases.""" + + def __init__(self, client: QdrantClient, embedding_dim: int): + self.client = client + self.embedding_dim = embedding_dim + + def ensure_collection_exists(self, collection_name: str) -> None: + """Ensure Qdrant collection exists with proper configuration. + + If the collection_name is actually an alias, use the aliased collection instead. + """ + if self.alias_exists(collection_name): + logger.info(f"Collection name {collection_name} is an alias, using existing aliased collection") + return + + collections = self.client.get_collections().collections + collection_names = [c.name for c in collections] + logger.debug(f"Existing collections: {collection_names}") + + if collection_name not in collection_names: + logger.info(f"Creating Qdrant collection: {collection_name}") + self.client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=self.embedding_dim, + distance=Distance.COSINE + ) + ) + logger.info(f"Created collection {collection_name}") + self._ensure_payload_indexes(collection_name) + else: + logger.info(f"Collection {collection_name} already exists") + + def create_versioned_collection(self, base_name: str) -> str: + """Create a new versioned collection for atomic swap indexing.""" + # Use milliseconds to avoid collisions in rapid calls + versioned_name = f"{base_name}_v{int(time.time() * 1000)}" + logger.info(f"Creating versioned collection: {versioned_name}") + + self.client.create_collection( + collection_name=versioned_name, + vectors_config=VectorParams( + size=self.embedding_dim, + distance=Distance.COSINE + ) + ) + self._ensure_payload_indexes(versioned_name) + return versioned_name + + def _ensure_payload_indexes(self, collection_name: str) -> None: + """Create payload indexes for efficient filtering on common fields.""" + try: + # Keyword index on 'path' for exact match and prefix filtering + self.client.create_payload_index( + collection_name=collection_name, + field_name="path", + field_schema=PayloadSchemaType.KEYWORD, + ) + # Keyword index on 'branch' for branch filtering + self.client.create_payload_index( + collection_name=collection_name, + field_name="branch", + field_schema=PayloadSchemaType.KEYWORD, + ) + logger.info(f"Payload indexes created for {collection_name}") + except Exception as e: + logger.warning(f"Failed to create payload indexes for {collection_name}: {e}") + + def delete_collection(self, collection_name: str) -> bool: + """Delete a collection.""" + try: + self.client.delete_collection(collection_name) + logger.info(f"Deleted collection: {collection_name}") + return True + except Exception as e: + logger.warning(f"Failed to delete collection {collection_name}: {e}") + return False + + def collection_exists(self, collection_name: str) -> bool: + """Check if a collection exists (not alias).""" + collections = self.client.get_collections().collections + return collection_name in [c.name for c in collections] + + def get_collection_names(self) -> List[str]: + """Get all collection names.""" + collections = self.client.get_collections().collections + return [c.name for c in collections] + + # Alias operations + + def alias_exists(self, alias_name: str) -> bool: + """Check if an alias exists.""" + try: + aliases = self.client.get_aliases() + exists = any(a.alias_name == alias_name for a in aliases.aliases) + logger.debug(f"Checking if alias '{alias_name}' exists: {exists}") + return exists + except Exception as e: + logger.warning(f"Error checking alias {alias_name}: {e}") + return False + + def resolve_alias(self, alias_name: str) -> Optional[str]: + """Resolve an alias to its underlying collection name.""" + try: + aliases = self.client.get_aliases() + for alias in aliases.aliases: + if alias.alias_name == alias_name: + return alias.collection_name + except Exception as e: + logger.debug(f"Error resolving alias {alias_name}: {e}") + return None + + def atomic_alias_swap( + self, + alias_name: str, + new_collection: str, + old_alias_exists: bool + ) -> None: + """Perform atomic alias swap for zero-downtime reindexing.""" + alias_operations = [] + + if old_alias_exists: + alias_operations.append( + DeleteAliasOperation(delete_alias=DeleteAlias(alias_name=alias_name)) + ) + + alias_operations.append( + CreateAliasOperation(create_alias=CreateAlias( + alias_name=alias_name, + collection_name=new_collection + )) + ) + + self.client.update_collection_aliases( + change_aliases_operations=alias_operations + ) + logger.info(f"Alias swap completed: {alias_name} -> {new_collection}") + + def delete_alias(self, alias_name: str) -> bool: + """Delete an alias.""" + try: + self.client.delete_alias(alias_name) + logger.info(f"Deleted alias: {alias_name}") + return True + except Exception as e: + logger.warning(f"Failed to delete alias {alias_name}: {e}") + return False + + def cleanup_orphaned_versioned_collections( + self, + base_name: str, + current_target: Optional[str] = None, + exclude_name: Optional[str] = None + ) -> int: + """Clean up orphaned versioned collections from failed indexing attempts.""" + cleaned = 0 + collection_names = self.get_collection_names() + + for coll_name in collection_names: + if coll_name.startswith(f"{base_name}_v") and coll_name != exclude_name: + if current_target != coll_name: + logger.info(f"Cleaning up orphaned versioned collection: {coll_name}") + if self.delete_collection(coll_name): + cleaned += 1 + + return cleaned diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/indexer.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/indexer.py new file mode 100644 index 00000000..259ccea5 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/indexer.py @@ -0,0 +1,393 @@ +""" +Repository indexing operations. + +Handles full repository indexing with atomic swap and streaming processing. +""" + +import gc +import logging +import random +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional, List + +from qdrant_client.models import Filter, FieldCondition, MatchAny, MatchValue + +from ...models.config import RAGConfig, IndexStats +from ...utils.utils import make_namespace +from .collection_manager import CollectionManager +from .branch_manager import BranchManager +from .point_operations import PointOperations +from .stats_manager import StatsManager + +logger = logging.getLogger(__name__) + +# Memory-efficient batch sizes +DOCUMENT_BATCH_SIZE = 50 +INSERT_BATCH_SIZE = 50 + + +class RepositoryIndexer: + """Handles repository indexing operations.""" + + def __init__( + self, + config: RAGConfig, + collection_manager: CollectionManager, + branch_manager: BranchManager, + point_ops: PointOperations, + stats_manager: StatsManager, + splitter, + loader + ): + self.config = config + self.collection_manager = collection_manager + self.branch_manager = branch_manager + self.point_ops = point_ops + self.stats_manager = stats_manager + self.splitter = splitter + self.loader = loader + + def estimate_repository_size( + self, + repo_path: str, + exclude_patterns: Optional[List[str]] = None + ) -> tuple[int, int]: + """Estimate repository size (file count and chunk count) without actually indexing.""" + logger.info(f"Estimating repository size for: {repo_path}") + + repo_path_obj = Path(repo_path) + file_list = list(self.loader.iter_repository_files(repo_path_obj, exclude_patterns)) + file_count = len(file_list) + logger.info(f"Found {file_count} files for estimation") + + if file_count == 0: + return 0, 0 + + SAMPLE_SIZE = 100 + chunk_count = 0 + + if file_count <= SAMPLE_SIZE: + for i in range(0, file_count, DOCUMENT_BATCH_SIZE): + batch = file_list[i:i + DOCUMENT_BATCH_SIZE] + documents = self.loader.load_file_batch( + batch, repo_path_obj, "estimate", "estimate", "estimate", "estimate" + ) + if documents: + chunks = self.splitter.split_documents(documents) + chunk_count += len(chunks) + del chunks + del documents + gc.collect() + else: + sample_files = random.sample(file_list, SAMPLE_SIZE) + sample_chunk_count = 0 + + for i in range(0, len(sample_files), DOCUMENT_BATCH_SIZE): + batch = sample_files[i:i + DOCUMENT_BATCH_SIZE] + documents = self.loader.load_file_batch( + batch, repo_path_obj, "estimate", "estimate", "estimate", "estimate" + ) + if documents: + chunks = self.splitter.split_documents(documents) + sample_chunk_count += len(chunks) + del chunks + del documents + gc.collect() + + avg_chunks_per_file = sample_chunk_count / SAMPLE_SIZE + chunk_count = int(avg_chunks_per_file * file_count) + logger.info(f"Estimated ~{avg_chunks_per_file:.1f} chunks/file from {SAMPLE_SIZE} samples") + gc.collect() + + logger.info(f"Estimated {chunk_count} chunks from {file_count} files") + return file_count, chunk_count + + def index_repository( + self, + repo_path: str, + workspace: str, + project: str, + branch: str, + commit: str, + alias_name: str, + exclude_patterns: Optional[List[str]] = None + ) -> IndexStats: + """Index entire repository for a branch using atomic swap strategy.""" + logger.info(f"Indexing repository: {workspace}/{project}/{branch} from {repo_path}") + + repo_path_obj = Path(repo_path) + temp_collection_name = self.collection_manager.create_versioned_collection(alias_name) + + # Check existing collection and preserve other branch data using streaming + old_collection_exists = self.collection_manager.alias_exists(alias_name) + if not old_collection_exists: + old_collection_exists = self.collection_manager.collection_exists(alias_name) + + actual_old_collection = None + if old_collection_exists: + actual_old_collection = self.collection_manager.resolve_alias(alias_name) or alias_name + + # Clean up orphaned versioned collections + current_target = self.collection_manager.resolve_alias(alias_name) + self.collection_manager.cleanup_orphaned_versioned_collections( + alias_name, current_target, temp_collection_name + ) + + # Get file list + file_list = list(self.loader.iter_repository_files(repo_path_obj, exclude_patterns)) + total_files = len(file_list) + logger.info(f"Found {total_files} files to index for branch '{branch}'") + + if total_files == 0: + logger.warning("No documents to index") + self.collection_manager.delete_collection(temp_collection_name) + return self.stats_manager.get_branch_stats( + workspace, project, branch, + self.collection_manager.resolve_alias(alias_name) or alias_name + ) + + # Validate limits + if self.config.max_files_per_index > 0 and total_files > self.config.max_files_per_index: + self.collection_manager.delete_collection(temp_collection_name) + raise ValueError( + f"Repository exceeds file limit: {total_files} files (max: {self.config.max_files_per_index})." + ) + + if self.config.max_chunks_per_index > 0: + logger.info("Estimating chunk count before indexing...") + _, estimated_chunks = self.estimate_repository_size(repo_path, exclude_patterns) + if estimated_chunks > self.config.max_chunks_per_index * 1.2: + self.collection_manager.delete_collection(temp_collection_name) + raise ValueError( + f"Repository estimated to exceed chunk limit: ~{estimated_chunks} chunks (max: {self.config.max_chunks_per_index})." + ) + + document_count = 0 + chunk_count = 0 + successful_chunks = 0 + failed_chunks = 0 + + try: + # Stream copy preserved points from other branches (memory-efficient) + if actual_old_collection: + self.branch_manager.stream_copy_points_to_collection( + actual_old_collection, + temp_collection_name, + branch, + INSERT_BATCH_SIZE + ) + + # Stream process files in batches + logger.info("Starting memory-efficient streaming indexing...") + batch_num = 0 + total_batches = (total_files + DOCUMENT_BATCH_SIZE - 1) // DOCUMENT_BATCH_SIZE + + for i in range(0, total_files, DOCUMENT_BATCH_SIZE): + batch_num += 1 + file_batch = file_list[i:i + DOCUMENT_BATCH_SIZE] + + documents = self.loader.load_file_batch( + file_batch, repo_path_obj, workspace, project, branch, commit + ) + document_count += len(documents) + + if not documents: + continue + + chunks = self.splitter.split_documents(documents) + batch_chunk_count = len(chunks) + chunk_count += batch_chunk_count + + # Check chunk limit + if self.config.max_chunks_per_index > 0 and chunk_count > self.config.max_chunks_per_index: + self.collection_manager.delete_collection(temp_collection_name) + raise ValueError(f"Repository exceeds chunk limit: {chunk_count}+ chunks.") + + # Process and upsert + success, failed = self.point_ops.process_and_upsert_chunks( + chunks, temp_collection_name, workspace, project, branch + ) + successful_chunks += success + failed_chunks += failed + + logger.info( + f"Batch {batch_num}/{total_batches}: processed {len(documents)} files, " + f"{batch_chunk_count} chunks" + ) + + del documents + del chunks + + if batch_num % 5 == 0: + gc.collect() + + logger.info( + f"Streaming indexing complete: {document_count} files, " + f"{successful_chunks}/{chunk_count} chunks indexed ({failed_chunks} failed)" + ) + + # Verify and perform atomic swap + temp_info = self.point_ops.client.get_collection(temp_collection_name) + if temp_info.points_count == 0: + raise Exception("Temporary collection is empty after indexing") + + self._perform_atomic_swap( + alias_name, temp_collection_name, old_collection_exists + ) + + except Exception as e: + logger.error(f"Indexing failed: {e}") + self.collection_manager.delete_collection(temp_collection_name) + raise e + finally: + gc.collect() + + self.stats_manager.store_metadata( + workspace, project, branch, commit, document_count, chunk_count + ) + + namespace = make_namespace(workspace, project, branch) + return IndexStats( + namespace=namespace, + document_count=document_count, + chunk_count=successful_chunks, + last_updated=datetime.now(timezone.utc).isoformat(), + workspace=workspace, + project=project, + branch=branch + ) + + def _perform_atomic_swap( + self, + alias_name: str, + temp_collection_name: str, + old_collection_exists: bool + ) -> None: + """Perform atomic alias swap with migration handling.""" + logger.info("Performing atomic alias swap...") + + is_direct_collection = ( + self.collection_manager.collection_exists(alias_name) and + not self.collection_manager.alias_exists(alias_name) + ) + + old_versioned_name = None + if old_collection_exists and not is_direct_collection: + old_versioned_name = self.collection_manager.resolve_alias(alias_name) + + try: + self.collection_manager.atomic_alias_swap( + alias_name, temp_collection_name, + old_collection_exists and not is_direct_collection + ) + except Exception as alias_err: + if is_direct_collection and "already exists" in str(alias_err).lower(): + logger.info("Migrating from direct collection to alias-based indexing...") + self.collection_manager.delete_collection(alias_name) + self.collection_manager.atomic_alias_swap(alias_name, temp_collection_name, False) + else: + raise alias_err + + if old_versioned_name and old_versioned_name != temp_collection_name: + self.collection_manager.delete_collection(old_versioned_name) + + +class FileOperations: + """Handles individual file update and delete operations.""" + + def __init__( + self, + client, + point_ops: PointOperations, + collection_manager: CollectionManager, + stats_manager: StatsManager, + splitter, + loader + ): + self.client = client + self.point_ops = point_ops + self.collection_manager = collection_manager + self.stats_manager = stats_manager + self.splitter = splitter + self.loader = loader + + def update_files( + self, + file_paths: List[str], + repo_base: str, + workspace: str, + project: str, + branch: str, + commit: str, + collection_name: str + ) -> IndexStats: + """Update specific files in the index (Delete Old -> Insert New).""" + logger.info(f"Updating {len(file_paths)} files in {workspace}/{project} for branch '{branch}'") + + repo_base_obj = Path(repo_base) + file_path_objs = [Path(fp) for fp in file_paths] + + self.collection_manager.ensure_collection_exists(collection_name) + + # Delete old chunks for these files and branch + logger.info(f"Purging existing vectors for {len(file_paths)} files in branch '{branch}'...") + self.client.delete( + collection_name=collection_name, + points_selector=Filter( + must=[ + FieldCondition(key="path", match=MatchAny(any=file_paths)), + FieldCondition(key="branch", match=MatchValue(value=branch)) + ] + ) + ) + + # Load and split new content + documents = self.loader.load_specific_files( + file_paths=file_path_objs, + repo_base=repo_base_obj, + workspace=workspace, + project=project, + branch=branch, + commit=commit + ) + + if not documents: + logger.warning("No documents loaded from provided paths.") + return self.stats_manager.get_project_stats(workspace, project, collection_name) + + chunks = self.splitter.split_documents(documents) + logger.info(f"Generated {len(chunks)} new chunks") + + # Process and upsert + self.point_ops.process_and_upsert_chunks( + chunks, collection_name, workspace, project, branch + ) + + logger.info(f"Successfully updated {len(chunks)} chunks for branch '{branch}'") + return self.stats_manager.get_project_stats(workspace, project, collection_name) + + def delete_files( + self, + file_paths: List[str], + workspace: str, + project: str, + branch: str, + collection_name: str + ) -> IndexStats: + """Delete specific files from the index for a specific branch.""" + logger.info(f"Deleting {len(file_paths)} files from {workspace}/{project} branch '{branch}'") + + self.client.delete( + collection_name=collection_name, + points_selector=Filter( + must=[ + FieldCondition(key="path", match=MatchAny(any=file_paths)), + FieldCondition(key="branch", match=MatchValue(value=branch)) + ] + ) + ) + logger.info(f"Deleted {len(file_paths)} files from branch '{branch}'") + + return self.stats_manager.get_project_stats(workspace, project, collection_name) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/manager.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/manager.py new file mode 100644 index 00000000..abd1d174 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/manager.py @@ -0,0 +1,282 @@ +""" +Main RAG Index Manager. + +Composes all index management components and provides the public API. +""" + +import logging +from typing import Optional, List + +from llama_index.core import Settings +from qdrant_client import QdrantClient + +from ...models.config import RAGConfig, IndexStats +from ...utils.utils import make_namespace, make_project_namespace +from ..splitter import ASTCodeSplitter +from ..loader import DocumentLoader +from ..embedding_factory import create_embedding_model, get_embedding_model_info + +from .collection_manager import CollectionManager +from .branch_manager import BranchManager +from .point_operations import PointOperations +from .stats_manager import StatsManager +from .indexer import RepositoryIndexer, FileOperations + +logger = logging.getLogger(__name__) + + +class RAGIndexManager: + """Manage RAG indices for code repositories using Qdrant. + + This is the main entry point for all indexing operations. + """ + + def __init__(self, config: RAGConfig): + self.config = config + + # Qdrant client + self.qdrant_client = QdrantClient(url=config.qdrant_url) + logger.info(f"Connected to Qdrant at {config.qdrant_url}") + + # Embedding model (supports Ollama and OpenRouter via factory) + embed_info = get_embedding_model_info(config) + logger.info(f"Using embedding provider: {embed_info['provider']} ({embed_info['type']})") + logger.info(f"Embedding model: {embed_info['model']}, dimension: {embed_info['embedding_dim']}") + + self.embed_model = create_embedding_model(config) + + # Global settings + Settings.embed_model = self.embed_model + Settings.chunk_size = config.chunk_size + Settings.chunk_overlap = config.chunk_overlap + + # Splitter and loader + logger.info("Using ASTCodeSplitter for code chunking (tree-sitter query-based)") + self.splitter = ASTCodeSplitter( + max_chunk_size=config.chunk_size, + min_chunk_size=min(200, config.chunk_size // 4), + chunk_overlap=config.chunk_overlap, + parser_threshold=10 + ) + self.loader = DocumentLoader(config) + + # Component managers + self._collection_manager = CollectionManager( + self.qdrant_client, config.embedding_dim + ) + self._branch_manager = BranchManager(self.qdrant_client) + self._point_ops = PointOperations( + self.qdrant_client, self.embed_model, batch_size=50 + ) + self._stats_manager = StatsManager( + self.qdrant_client, config.qdrant_collection_prefix + ) + + # Higher-level operations + self._indexer = RepositoryIndexer( + config=config, + collection_manager=self._collection_manager, + branch_manager=self._branch_manager, + point_ops=self._point_ops, + stats_manager=self._stats_manager, + splitter=self.splitter, + loader=self.loader + ) + self._file_ops = FileOperations( + client=self.qdrant_client, + point_ops=self._point_ops, + collection_manager=self._collection_manager, + stats_manager=self._stats_manager, + splitter=self.splitter, + loader=self.loader + ) + + # Collection naming + + def _get_project_collection_name(self, workspace: str, project: str) -> str: + """Generate Qdrant collection name from workspace/project.""" + namespace = make_project_namespace(workspace, project) + return f"{self.config.qdrant_collection_prefix}_{namespace}" + + # Repository indexing + + def estimate_repository_size( + self, + repo_path: str, + exclude_patterns: Optional[List[str]] = None + ) -> tuple[int, int]: + """Estimate repository size (file count and chunk count).""" + return self._indexer.estimate_repository_size(repo_path, exclude_patterns) + + def index_repository( + self, + repo_path: str, + workspace: str, + project: str, + branch: str, + commit: str, + exclude_patterns: Optional[List[str]] = None + ) -> IndexStats: + """Index entire repository for a branch using atomic swap strategy.""" + alias_name = self._get_project_collection_name(workspace, project) + return self._indexer.index_repository( + repo_path=repo_path, + workspace=workspace, + project=project, + branch=branch, + commit=commit, + alias_name=alias_name, + exclude_patterns=exclude_patterns + ) + + # File operations + + def update_files( + self, + file_paths: List[str], + repo_base: str, + workspace: str, + project: str, + branch: str, + commit: str + ) -> IndexStats: + """Update specific files in the index (Delete Old -> Insert New).""" + collection_name = self._get_project_collection_name(workspace, project) + return self._file_ops.update_files( + file_paths=file_paths, + repo_base=repo_base, + workspace=workspace, + project=project, + branch=branch, + commit=commit, + collection_name=collection_name + ) + + def delete_files( + self, + file_paths: List[str], + workspace: str, + project: str, + branch: str + ) -> IndexStats: + """Delete specific files from the index for a specific branch.""" + collection_name = self._get_project_collection_name(workspace, project) + return self._file_ops.delete_files( + file_paths=file_paths, + workspace=workspace, + project=project, + branch=branch, + collection_name=collection_name + ) + + # Branch operations + + def delete_branch(self, workspace: str, project: str, branch: str) -> bool: + """Delete all points for a specific branch from the project collection.""" + collection_name = self._get_project_collection_name(workspace, project) + + if not self._collection_manager.collection_exists(collection_name): + if not self._collection_manager.alias_exists(collection_name): + logger.warning(f"Collection {collection_name} does not exist") + return False + + return self._branch_manager.delete_branch_points(collection_name, branch) + + def get_branch_point_count(self, workspace: str, project: str, branch: str) -> int: + """Get the number of points for a specific branch.""" + collection_name = self._get_project_collection_name(workspace, project) + + if not self._collection_manager.collection_exists(collection_name): + if not self._collection_manager.alias_exists(collection_name): + return 0 + + return self._branch_manager.get_branch_point_count(collection_name, branch) + + def get_indexed_branches(self, workspace: str, project: str) -> List[str]: + """Get list of branches that have points in the collection.""" + collection_name = self._get_project_collection_name(workspace, project) + + if not self._collection_manager.collection_exists(collection_name): + if not self._collection_manager.alias_exists(collection_name): + return [] + + return self._branch_manager.get_indexed_branches(collection_name) + + # Index management + + def delete_index(self, workspace: str, project: str, branch: str): + """Delete branch data from project index.""" + if branch and branch != "*": + self.delete_branch(workspace, project, branch) + else: + self.delete_project_index(workspace, project) + + def delete_project_index(self, workspace: str, project: str): + """Delete entire project collection (all branches).""" + collection_name = self._get_project_collection_name(workspace, project) + namespace = make_project_namespace(workspace, project) + + logger.info(f"Deleting entire project index for {namespace}") + + try: + if self._collection_manager.alias_exists(collection_name): + actual_collection = self._collection_manager.resolve_alias(collection_name) + self._collection_manager.delete_alias(collection_name) + if actual_collection: + self._collection_manager.delete_collection(actual_collection) + else: + self._collection_manager.delete_collection(collection_name) + logger.info(f"Deleted Qdrant collection: {collection_name}") + except Exception as e: + logger.warning(f"Failed to delete Qdrant collection: {e}") + + # Statistics + + def _get_index_stats(self, workspace: str, project: str, branch: str) -> IndexStats: + """Get statistics about a branch index (backward compatibility).""" + return self._get_branch_index_stats(workspace, project, branch) + + def _get_branch_index_stats(self, workspace: str, project: str, branch: str) -> IndexStats: + """Get statistics about a specific branch within a project collection.""" + collection_name = self._get_project_collection_name(workspace, project) + return self._stats_manager.get_branch_stats( + workspace, project, branch, collection_name + ) + + def _get_project_index_stats(self, workspace: str, project: str) -> IndexStats: + """Get statistics about a project's index (all branches combined).""" + collection_name = self._get_project_collection_name(workspace, project) + return self._stats_manager.get_project_stats(workspace, project, collection_name) + + def list_indices(self) -> List[IndexStats]: + """List all project indices with branch breakdown.""" + return self._stats_manager.list_all_indices( + self._collection_manager.alias_exists + ) + + # Legacy/compatibility methods + + def _ensure_collection_exists(self, collection_name: str): + """Ensure Qdrant collection exists (legacy compatibility).""" + self._collection_manager.ensure_collection_exists(collection_name) + + def _alias_exists(self, alias_name: str) -> bool: + """Check if an alias exists (legacy compatibility).""" + return self._collection_manager.alias_exists(alias_name) + + def _resolve_alias_to_collection(self, alias_name: str) -> Optional[str]: + """Resolve an alias to its collection (legacy compatibility).""" + return self._collection_manager.resolve_alias(alias_name) + + def _generate_point_id( + self, + workspace: str, + project: str, + branch: str, + path: str, + chunk_index: int + ) -> str: + """Generate deterministic point ID (legacy compatibility).""" + return PointOperations.generate_point_id( + workspace, project, branch, path, chunk_index + ) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/point_operations.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/point_operations.py new file mode 100644 index 00000000..409e8462 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/point_operations.py @@ -0,0 +1,155 @@ +""" +Point operations for embedding and upserting vectors. + +Handles embedding generation, point creation, and batch upsert operations. +""" + +import logging +import uuid +from datetime import datetime, timezone +from typing import List, Dict, Tuple + +from llama_index.core.schema import TextNode +from qdrant_client import QdrantClient +from qdrant_client.models import PointStruct + +logger = logging.getLogger(__name__) + + +class PointOperations: + """Handles point embedding and upsert operations.""" + + def __init__(self, client: QdrantClient, embed_model, batch_size: int = 50): + self.client = client + self.embed_model = embed_model + self.batch_size = batch_size + + @staticmethod + def generate_point_id( + workspace: str, + project: str, + branch: str, + path: str, + chunk_index: int + ) -> str: + """Generate deterministic point ID for upsert (same content = same ID = replace).""" + key = f"{workspace}:{project}:{branch}:{path}:{chunk_index}" + return str(uuid.uuid5(uuid.NAMESPACE_DNS, key)) + + def prepare_chunks_for_embedding( + self, + chunks: List[TextNode], + workspace: str, + project: str, + branch: str + ) -> List[Tuple[str, TextNode]]: + """Prepare chunks with deterministic IDs for embedding. + + Returns list of (point_id, chunk) tuples. + """ + # Group chunks by file path + chunks_by_file: Dict[str, List[TextNode]] = {} + for chunk in chunks: + path = chunk.metadata.get("path") + if not path: + # Log warning - missing path breaks idempotency but shouldn't happen in normal flow + logger.warning("Chunk missing 'path' metadata - using non-deterministic ID (breaks idempotency)") + path = f"__unknown__{uuid.uuid4()}" + if path not in chunks_by_file: + chunks_by_file[path] = [] + chunks_by_file[path].append(chunk) + + # Assign deterministic IDs + chunk_data = [] + for path, file_chunks in chunks_by_file.items(): + for chunk_index, chunk in enumerate(file_chunks): + point_id = self.generate_point_id(workspace, project, branch, path, chunk_index) + chunk.metadata["indexed_at"] = datetime.now(timezone.utc).isoformat() + chunk_data.append((point_id, chunk)) + + return chunk_data + + def embed_and_create_points( + self, + chunk_data: List[Tuple[str, TextNode]] + ) -> List[PointStruct]: + """Embed chunks and create Qdrant points. + + Args: + chunk_data: List of (point_id, chunk) tuples + + Returns: + List of PointStruct ready for upsert + """ + if not chunk_data: + return [] + + # Batch embed all chunks at once + texts_to_embed = [chunk.text for _, chunk in chunk_data] + embeddings = self.embed_model.get_text_embedding_batch(texts_to_embed) + + # Build points with embeddings + points = [] + for (point_id, chunk), embedding in zip(chunk_data, embeddings): + points.append(PointStruct( + id=point_id, + vector=embedding, + payload={ + **chunk.metadata, + "text": chunk.text, + "_node_content": chunk.text, + } + )) + + return points + + def upsert_points( + self, + collection_name: str, + points: List[PointStruct] + ) -> Tuple[int, int]: + """Upsert points to collection in batches. + + Returns: + Tuple of (successful_count, failed_count) + """ + successful = 0 + failed = 0 + + for i in range(0, len(points), self.batch_size): + batch = points[i:i + self.batch_size] + try: + self.client.upsert( + collection_name=collection_name, + points=batch + ) + successful += len(batch) + except Exception as e: + logger.error(f"Failed to upsert batch starting at {i}: {e}") + failed += len(batch) + + return successful, failed + + def process_and_upsert_chunks( + self, + chunks: List[TextNode], + collection_name: str, + workspace: str, + project: str, + branch: str + ) -> Tuple[int, int]: + """Full pipeline: prepare, embed, and upsert chunks. + + Returns: + Tuple of (successful_count, failed_count) + """ + # Prepare chunks with IDs + chunk_data = self.prepare_chunks_for_embedding( + chunks, workspace, project, branch + ) + + # Embed and create points + points = self.embed_and_create_points(chunk_data) + + # Upsert to collection + return self.upsert_points(collection_name, points) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/stats_manager.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/stats_manager.py new file mode 100644 index 00000000..226f021f --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/index_manager/stats_manager.py @@ -0,0 +1,156 @@ +""" +Index statistics and metadata operations. +""" + +import logging +from datetime import datetime, timezone +from typing import List, Optional + +from qdrant_client import QdrantClient +from qdrant_client.models import Filter, FieldCondition, MatchValue + +from ...models.config import IndexStats +from ...utils.utils import make_namespace, make_project_namespace + +logger = logging.getLogger(__name__) + + +class StatsManager: + """Manages index statistics and metadata.""" + + def __init__(self, client: QdrantClient, collection_prefix: str): + self.client = client + self.collection_prefix = collection_prefix + + def get_branch_stats( + self, + workspace: str, + project: str, + branch: str, + collection_name: str + ) -> IndexStats: + """Get statistics about a specific branch within a project collection.""" + namespace = make_namespace(workspace, project, branch) + + try: + count_result = self.client.count( + collection_name=collection_name, + count_filter=Filter( + must=[ + FieldCondition( + key="branch", + match=MatchValue(value=branch) + ) + ] + ) + ) + chunk_count = count_result.count + + return IndexStats( + namespace=namespace, + document_count=0, + chunk_count=chunk_count, + last_updated=datetime.now(timezone.utc).isoformat(), + workspace=workspace, + project=project, + branch=branch + ) + except Exception: + return IndexStats( + namespace=namespace, + document_count=0, + chunk_count=0, + last_updated="", + workspace=workspace, + project=project, + branch=branch + ) + + def get_project_stats( + self, + workspace: str, + project: str, + collection_name: str + ) -> IndexStats: + """Get statistics about a project's index (all branches combined).""" + namespace = make_project_namespace(workspace, project) + + try: + collection_info = self.client.get_collection(collection_name) + chunk_count = collection_info.points_count + + return IndexStats( + namespace=namespace, + document_count=0, + chunk_count=chunk_count, + last_updated=datetime.now(timezone.utc).isoformat(), + workspace=workspace, + project=project, + branch="*" + ) + except Exception: + return IndexStats( + namespace=namespace, + document_count=0, + chunk_count=0, + last_updated="", + workspace=workspace, + project=project, + branch="*" + ) + + def list_all_indices(self, alias_checker) -> List[IndexStats]: + """List all project indices with branch breakdown. + + Args: + alias_checker: Function to check if name is an alias + """ + indices = [] + collections = self.client.get_collections().collections + + for collection in collections: + if collection.name.startswith(f"{self.collection_prefix}_"): + namespace = collection.name[len(f"{self.collection_prefix}_"):] + parts = namespace.split("__") + + if len(parts) == 2: + # New format: workspace__project + workspace, project = parts + stats = self.get_project_stats( + workspace, project, collection.name + ) + indices.append(stats) + elif len(parts) == 3: + # Legacy format: workspace__project__branch + workspace, project, branch = parts + stats = self.get_branch_stats( + workspace, project, branch, collection.name + ) + indices.append(stats) + + return indices + + def store_metadata( + self, + workspace: str, + project: str, + branch: str, + commit: str, + document_count: int, + chunk_count: int + ) -> None: + """Store/log metadata for an indexing operation.""" + namespace = make_namespace(workspace, project, branch) + + metadata = { + "namespace": namespace, + "workspace": workspace, + "project": project, + "branch": branch, + "commit": commit, + "document_count": document_count, + "chunk_count": chunk_count, + "last_updated": datetime.now(timezone.utc).isoformat(), + } + + logger.info(f"Indexed {namespace}: {document_count} docs, {chunk_count} chunks") diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/loader.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/loader.py index 25d92966..05235278 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/loader.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/loader.py @@ -3,7 +3,7 @@ import logging from llama_index.core.schema import Document -from ..utils.utils import detect_language_from_path, should_exclude_file, is_binary_file +from ..utils.utils import detect_language_from_path, should_exclude_file, is_binary_file, clean_archive_path from ..models.config import RAGConfig logger = logging.getLogger(__name__) @@ -106,11 +106,14 @@ def load_file_batch( language = detect_language_from_path(str(full_path)) filetype = full_path.suffix.lstrip('.') + # Clean archive root prefix from path (e.g., 'owner-repo-commit/src/file.php' -> 'src/file.php') + clean_path = clean_archive_path(relative_path_str) + metadata = { "workspace": workspace, "project": project, "branch": branch, - "path": relative_path_str, + "path": clean_path, "commit": commit, "language": language, "filetype": filetype, @@ -190,11 +193,14 @@ def load_from_directory( language = detect_language_from_path(str(file_path)) filetype = file_path.suffix.lstrip('.') + # Clean archive root prefix from path + clean_path = clean_archive_path(relative_path) + metadata = { "workspace": workspace, "project": project, "branch": branch, - "path": relative_path, + "path": clean_path, "commit": commit, "language": language, "filetype": filetype, @@ -207,7 +213,7 @@ def load_from_directory( ) documents.append(doc) - logger.debug(f"Loaded document: {relative_path} ({language})") + logger.debug(f"Loaded document: {clean_path} ({language})") logger.info(f"Loaded {len(documents)} documents from {repo_path} (excluded {excluded_count} files by patterns)") return documents @@ -257,11 +263,14 @@ def load_specific_files( language = detect_language_from_path(str(full_path)) filetype = full_path.suffix.lstrip('.') + # Clean archive root prefix from path + clean_path = clean_archive_path(relative_path) + metadata = { "workspace": workspace, "project": project, "branch": branch, - "path": relative_path, + "path": clean_path, "commit": commit, "language": language, "filetype": filetype, @@ -274,7 +283,7 @@ def load_specific_files( ) documents.append(doc) - logger.debug(f"Loaded document: {relative_path}") + logger.debug(f"Loaded document: {clean_path}") return documents diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/ollama_embedding.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/ollama_embedding.py new file mode 100644 index 00000000..5155a2f8 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/ollama_embedding.py @@ -0,0 +1,252 @@ +""" +Ollama embedding wrapper for LlamaIndex. +Supports local embedding models running via Ollama. +""" + +import os +from typing import Any, List, Optional +from llama_index.core.base.embeddings.base import BaseEmbedding +import httpx +import logging + +from ..models.config import get_embedding_dim_for_model + +logger = logging.getLogger(__name__) + +# Default batch size - can be overridden via OLLAMA_BATCH_SIZE env var +DEFAULT_BATCH_SIZE = int(os.getenv("OLLAMA_BATCH_SIZE", "100")) +# Default timeout - can be overridden via OLLAMA_TIMEOUT env var +DEFAULT_TIMEOUT = float(os.getenv("OLLAMA_TIMEOUT", "120")) + + +class OllamaEmbedding(BaseEmbedding): + """ + Custom embedding class for Ollama API. + + Supports local embedding models like qwen3-embedding:0.6b, nomic-embed-text, etc. + """ + + def __init__( + self, + model: str = "qwen3-embedding:0.6b", + base_url: str = "http://localhost:11434", + timeout: float = None, + embed_batch_size: int = None, + expected_dim: Optional[int] = None, + **kwargs: Any + ): + # Use env-configured defaults if not specified + if timeout is None: + timeout = DEFAULT_TIMEOUT + if embed_batch_size is None: + embed_batch_size = DEFAULT_BATCH_SIZE + + super().__init__(embed_batch_size=embed_batch_size, **kwargs) + + # Determine expected embedding dimension + if expected_dim is not None: + embedding_dim = expected_dim + else: + embedding_dim = get_embedding_dim_for_model(model) + + logger.info(f"OllamaEmbedding: Initializing with model: {model}") + logger.info(f"OllamaEmbedding: Expected embedding dimension: {embedding_dim}") + logger.info(f"OllamaEmbedding: Base URL: {base_url}") + logger.info(f"OllamaEmbedding: Batch size: {embed_batch_size}") + + # Store config using object.__setattr__ to bypass Pydantic validation + object.__setattr__(self, '_config', { + "model": model, + "base_url": base_url.rstrip('/'), + "timeout": timeout, + "embed_batch_size": embed_batch_size, + "embedding_dim": embedding_dim + }) + + # Initialize HTTP client + object.__setattr__(self, '_client', httpx.Client( + base_url=base_url.rstrip('/'), + timeout=timeout + )) + + # Test connection + self._test_connection() + logger.info(f"Ollama embeddings initialized successfully") + + def _test_connection(self): + """Test connection to Ollama server.""" + try: + response = self._client.get("/api/tags") + if response.status_code == 200: + models = response.json().get("models", []) + model_names = [m.get("name", "") for m in models] + logger.info(f"Connected to Ollama. Available models: {model_names}") + + # Check if our model is available + model_name = self._config["model"] + if not any(model_name in name or name in model_name for name in model_names): + logger.warning(f"Model '{model_name}' may not be available. Pull it with: ollama pull {model_name}") + else: + logger.warning(f"Could not list Ollama models: {response.status_code}") + except Exception as e: + logger.warning(f"Could not connect to Ollama at {self._config['base_url']}: {e}") + logger.warning("Make sure Ollama is running: ollama serve") + + def close(self): + """Close the HTTP client and free resources.""" + try: + if hasattr(self, '_client') and self._client: + self._client.close() + logger.info("Ollama embedding client closed") + except Exception as e: + logger.warning(f"Error closing Ollama client: {e}") + + def __del__(self): + """Destructor to ensure client is closed.""" + self.close() + + @property + def model(self) -> str: + """Get the model name.""" + return self._config["model"] + + def _get_query_embedding(self, query: str) -> List[float]: + """Get embedding for a query text.""" + return self._get_embedding(query) + + def _get_text_embedding(self, text: str) -> List[float]: + """Get embedding for a text.""" + return self._get_embedding(text) + + def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]: + """ + Get embeddings for multiple texts using batch API. + Uses /api/embed which supports array input for batching. + """ + if not texts: + return [] + + expected_dim = self._config.get("embedding_dim", 1024) + batch_size = self._config.get("embed_batch_size", DEFAULT_BATCH_SIZE) + all_embeddings = [] + + # Process in batches + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + logger.debug(f"Embedding batch {i // batch_size + 1}: {len(batch)} texts") + + # Preprocess batch + processed_batch = [] + for text in batch: + if not text or not text.strip(): + processed_batch.append(" ") # Placeholder for empty + else: + # Truncate if too long + if len(text) > 24000: + text = text[:24000] + processed_batch.append(text.strip()) + + try: + # Use /api/embed with array input for batch processing + response = self._client.post( + "/api/embed", + json={ + "model": self._config["model"], + "input": processed_batch + } + ) + response.raise_for_status() + data = response.json() + + # /api/embed returns {"embeddings": [[...], [...], ...]} + if "embeddings" in data: + batch_embeddings = data["embeddings"] + all_embeddings.extend(batch_embeddings) + else: + logger.error(f"Unexpected batch response format: {list(data.keys())}") + # Fallback to zeros + all_embeddings.extend([[0.0] * expected_dim] * len(processed_batch)) + + except Exception as e: + logger.error(f"Batch embedding failed: {e}, falling back to single requests") + # Fallback to single embedding requests + for text in processed_batch: + try: + embedding = self._get_embedding(text) + all_embeddings.append(embedding) + except Exception: + all_embeddings.append([0.0] * expected_dim) + + return all_embeddings + + def _get_embedding(self, text: str) -> List[float]: + """Get embedding from Ollama API.""" + expected_dim = self._config.get("embedding_dim", 1024) + + try: + # Validate input + if not text or not text.strip(): + logger.warning("Empty text provided for embedding, using placeholder") + return [0.0] * expected_dim + + # Truncate if too long (Ollama typically handles ~8k tokens) + max_chars = 24000 + if len(text) > max_chars: + logger.warning(f"Text too long ({len(text)} chars), truncating to {max_chars}") + text = text[:max_chars] + + # Clean the text + text = text.strip() + if not text: + logger.warning("Text became empty after stripping") + return [0.0] * expected_dim + + # Call Ollama embeddings API + response = self._client.post( + "/api/embeddings", + json={ + "model": self._config["model"], + "prompt": text + } + ) + response.raise_for_status() + data = response.json() + + # Ollama returns embedding in 'embedding' field + if "embedding" in data: + embedding = data["embedding"] + elif "embeddings" in data and len(data["embeddings"]) > 0: + # Fallback for potential future API changes + embedding = data["embeddings"][0] + else: + logger.error(f"Unexpected response format from Ollama: {data}") + return [0.0] * expected_dim + + # Validate embedding dimensions + if len(embedding) != expected_dim: + logger.warning(f"Unexpected embedding dimension: {len(embedding)}, expected {expected_dim}") + # Update expected dim if this is consistently different + self._config["embedding_dim"] = len(embedding) + + return embedding + + except httpx.HTTPStatusError as e: + logger.error(f"HTTP error getting embedding from Ollama: {e}") + logger.error(f"Response: {e.response.text if e.response else 'No response'}") + raise + except Exception as e: + logger.error(f"Error getting embedding from Ollama: {e}") + logger.error(f"Text length: {len(text) if text else 0}, Text preview: {text[:100] if text else 'None'}...") + raise + + async def _aget_query_embedding(self, query: str) -> List[float]: + """Async get embedding for a query text.""" + return self._get_query_embedding(query) + + async def _aget_text_embedding(self, text: str) -> List[float]: + """Async get embedding for a text.""" + return self._get_text_embedding(text) + + async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]: + """Async batch get embeddings for multiple texts.""" + return self._get_text_embeddings(texts) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/semantic_splitter.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/semantic_splitter.py deleted file mode 100644 index 349cdf23..00000000 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/semantic_splitter.py +++ /dev/null @@ -1,455 +0,0 @@ -""" -Semantic Code Splitter - Intelligent code splitting using LangChain's language-aware splitters. - -This module provides smart code chunking that: -1. Uses LangChain's RecursiveCharacterTextSplitter with language-specific separators -2. Supports 25+ programming languages out of the box -3. Enriches metadata with semantic information (function names, imports, etc.) -4. Falls back gracefully for unsupported languages -""" - -import re -import hashlib -import logging -from typing import List, Dict, Any, Optional -from dataclasses import dataclass, field -from enum import Enum - -from langchain_text_splitters import RecursiveCharacterTextSplitter, Language -from llama_index.core.schema import Document, TextNode - -logger = logging.getLogger(__name__) - - -class ChunkType(Enum): - """Type of code chunk for semantic understanding""" - CLASS = "class" - FUNCTION = "function" - METHOD = "method" - INTERFACE = "interface" - MODULE = "module" - IMPORTS = "imports" - CONSTANTS = "constants" - DOCUMENTATION = "documentation" - CONFIG = "config" - MIXED = "mixed" - UNKNOWN = "unknown" - - -@dataclass -class CodeBlock: - """Represents a logical block of code""" - content: str - chunk_type: ChunkType - name: Optional[str] = None - parent_name: Optional[str] = None - start_line: int = 0 - end_line: int = 0 - imports: List[str] = field(default_factory=list) - docstring: Optional[str] = None - signature: Optional[str] = None - - -# Map internal language names to LangChain Language enum -LANGUAGE_MAP: Dict[str, Language] = { - 'python': Language.PYTHON, - 'java': Language.JAVA, - 'kotlin': Language.KOTLIN, - 'javascript': Language.JS, - 'typescript': Language.TS, - 'go': Language.GO, - 'rust': Language.RUST, - 'php': Language.PHP, - 'ruby': Language.RUBY, - 'scala': Language.SCALA, - 'swift': Language.SWIFT, - 'c': Language.C, - 'cpp': Language.CPP, - 'csharp': Language.CSHARP, - 'markdown': Language.MARKDOWN, - 'html': Language.HTML, - 'latex': Language.LATEX, - 'rst': Language.RST, - 'lua': Language.LUA, - 'perl': Language.PERL, - 'haskell': Language.HASKELL, - 'solidity': Language.SOL, - 'proto': Language.PROTO, - 'cobol': Language.COBOL, -} - -# Patterns for metadata extraction -METADATA_PATTERNS = { - 'python': { - 'class': re.compile(r'^class\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'^(?:async\s+)?def\s+(\w+)\s*\(', re.MULTILINE), - 'import': re.compile(r'^(?:from\s+[\w.]+\s+)?import\s+.+$', re.MULTILINE), - 'docstring': re.compile(r'"""[\s\S]*?"""|\'\'\'[\s\S]*?\'\'\''), - }, - 'java': { - 'class': re.compile(r'(?:public\s+|private\s+|protected\s+)?(?:abstract\s+|final\s+)?class\s+(\w+)', re.MULTILINE), - 'interface': re.compile(r'(?:public\s+)?interface\s+(\w+)', re.MULTILINE), - 'method': re.compile(r'(?:public|private|protected)\s+(?:static\s+)?(?:final\s+)?[\w<>,\s]+\s+(\w+)\s*\(', re.MULTILINE), - 'import': re.compile(r'^import\s+[\w.*]+;', re.MULTILINE), - }, - 'javascript': { - 'class': re.compile(r'(?:export\s+)?(?:default\s+)?class\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:export\s+)?(?:async\s+)?function\s*\*?\s*(\w+)\s*\(', re.MULTILINE), - 'arrow': re.compile(r'(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\([^)]*\)\s*=>', re.MULTILINE), - 'import': re.compile(r'^import\s+.*?from\s+[\'"]([^\'"]+)[\'"]', re.MULTILINE), - }, - 'typescript': { - 'class': re.compile(r'(?:export\s+)?(?:default\s+)?class\s+(\w+)', re.MULTILINE), - 'interface': re.compile(r'(?:export\s+)?interface\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:export\s+)?(?:async\s+)?function\s*\*?\s*(\w+)\s*\(', re.MULTILINE), - 'type': re.compile(r'(?:export\s+)?type\s+(\w+)', re.MULTILINE), - 'import': re.compile(r'^import\s+.*?from\s+[\'"]([^\'"]+)[\'"]', re.MULTILINE), - }, - 'go': { - 'function': re.compile(r'^func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(', re.MULTILINE), - 'struct': re.compile(r'^type\s+(\w+)\s+struct\s*\{', re.MULTILINE), - 'interface': re.compile(r'^type\s+(\w+)\s+interface\s*\{', re.MULTILINE), - }, - 'rust': { - 'function': re.compile(r'^(?:pub\s+)?(?:async\s+)?fn\s+(\w+)', re.MULTILINE), - 'struct': re.compile(r'^(?:pub\s+)?struct\s+(\w+)', re.MULTILINE), - 'impl': re.compile(r'^impl(?:<[^>]+>)?\s+(?:\w+\s+for\s+)?(\w+)', re.MULTILINE), - 'trait': re.compile(r'^(?:pub\s+)?trait\s+(\w+)', re.MULTILINE), - }, - 'php': { - 'class': re.compile(r'(?:abstract\s+|final\s+)?class\s+(\w+)', re.MULTILINE), - 'interface': re.compile(r'interface\s+(\w+)', re.MULTILINE), - 'function': re.compile(r'(?:public|private|protected|static|\s)*function\s+(\w+)\s*\(', re.MULTILINE), - }, - 'csharp': { - 'class': re.compile(r'(?:public\s+|private\s+|protected\s+)?(?:abstract\s+|sealed\s+)?class\s+(\w+)', re.MULTILINE), - 'interface': re.compile(r'(?:public\s+)?interface\s+(\w+)', re.MULTILINE), - 'method': re.compile(r'(?:public|private|protected)\s+(?:static\s+)?(?:async\s+)?[\w<>,\s]+\s+(\w+)\s*\(', re.MULTILINE), - }, -} - - -class SemanticCodeSplitter: - """ - Intelligent code splitter using LangChain's language-aware text splitters. - - Features: - - Uses LangChain's RecursiveCharacterTextSplitter with language-specific separators - - Supports 25+ programming languages (Python, Java, JS/TS, Go, Rust, PHP, etc.) - - Enriches chunks with semantic metadata (function names, classes, imports) - - Graceful fallback for unsupported languages - """ - - DEFAULT_CHUNK_SIZE = 1500 - DEFAULT_CHUNK_OVERLAP = 200 - DEFAULT_MIN_CHUNK_SIZE = 100 - - def __init__( - self, - max_chunk_size: int = DEFAULT_CHUNK_SIZE, - min_chunk_size: int = DEFAULT_MIN_CHUNK_SIZE, - overlap: int = DEFAULT_CHUNK_OVERLAP - ): - self.max_chunk_size = max_chunk_size - self.min_chunk_size = min_chunk_size - self.overlap = overlap - - # Cache splitters for reuse - self._splitter_cache: Dict[str, RecursiveCharacterTextSplitter] = {} - - # Default splitter for unknown languages - self._default_splitter = RecursiveCharacterTextSplitter( - chunk_size=max_chunk_size, - chunk_overlap=overlap, - length_function=len, - is_separator_regex=False, - ) - - @staticmethod - def _make_deterministic_id(namespace: str, path: str, chunk_index: int) -> str: - """Generate deterministic chunk ID for idempotent indexing""" - key = f"{namespace}:{path}:{chunk_index}" - return hashlib.sha256(key.encode()).hexdigest()[:32] - - def _get_splitter(self, language: str) -> RecursiveCharacterTextSplitter: - """Get or create a language-specific splitter""" - if language in self._splitter_cache: - return self._splitter_cache[language] - - lang_enum = LANGUAGE_MAP.get(language.lower()) - - if lang_enum: - splitter = RecursiveCharacterTextSplitter.from_language( - language=lang_enum, - chunk_size=self.max_chunk_size, - chunk_overlap=self.overlap, - ) - self._splitter_cache[language] = splitter - return splitter - - return self._default_splitter - - def split_documents(self, documents: List[Document]) -> List[TextNode]: - """Split documents into semantic chunks with enriched metadata""" - return list(self.iter_split_documents(documents)) - - def iter_split_documents(self, documents: List[Document]): - """Generator that yields chunks one at a time for memory efficiency""" - for doc in documents: - language = doc.metadata.get("language", "text") - path = doc.metadata.get("path", "unknown") - - try: - for node in self._split_document(doc, language): - yield node - except Exception as e: - logger.warning(f"Splitting failed for {path}: {e}, using fallback") - for node in self._fallback_split(doc): - yield node - - def _split_document(self, doc: Document, language: str) -> List[TextNode]: - """Split a single document using language-aware splitter""" - text = doc.text - - if not text or not text.strip(): - return [] - - # Get language-specific splitter - splitter = self._get_splitter(language) - - # Split the text - chunks = splitter.split_text(text) - - # Filter empty chunks and convert to nodes with metadata - nodes = [] - text_offset = 0 - - for i, chunk in enumerate(chunks): - if not chunk or not chunk.strip(): - continue - - # Skip very small chunks unless they're standalone - if len(chunk.strip()) < self.min_chunk_size and len(chunks) > 1: - # Try to find and merge with adjacent chunk - continue - - # Calculate approximate line numbers - start_line = text[:text_offset].count('\n') + 1 if text_offset > 0 else 1 - chunk_pos = text.find(chunk, text_offset) - if chunk_pos >= 0: - text_offset = chunk_pos + len(chunk) - end_line = start_line + chunk.count('\n') - - # Extract semantic metadata - metadata = self._extract_metadata(chunk, language, doc.metadata) - metadata.update({ - 'chunk_index': i, - 'total_chunks': len(chunks), - 'start_line': start_line, - 'end_line': end_line, - }) - - chunk_id = self._make_deterministic_id( - metadata.get('namespace', ''), - metadata.get('path', ''), - i - ) - node = TextNode( - id_=chunk_id, - text=chunk, - metadata=metadata - ) - nodes.append(node) - - return nodes - - def _extract_metadata( - self, - chunk: str, - language: str, - base_metadata: Dict[str, Any] - ) -> Dict[str, Any]: - """Extract semantic metadata from a code chunk""" - metadata = dict(base_metadata) - - # Determine chunk type and extract names - chunk_type = ChunkType.MIXED - names = [] - imports = [] - - patterns = METADATA_PATTERNS.get(language.lower(), {}) - - # Check for classes - if 'class' in patterns: - matches = patterns['class'].findall(chunk) - if matches: - chunk_type = ChunkType.CLASS - names.extend(matches) - - # Check for interfaces - if 'interface' in patterns: - matches = patterns['interface'].findall(chunk) - if matches: - chunk_type = ChunkType.INTERFACE - names.extend(matches) - - # Check for functions/methods - if chunk_type == ChunkType.MIXED: - for key in ['function', 'method', 'arrow']: - if key in patterns: - matches = patterns[key].findall(chunk) - if matches: - chunk_type = ChunkType.FUNCTION - names.extend(matches) - break - - # Check for imports - if 'import' in patterns: - import_matches = patterns['import'].findall(chunk) - if import_matches: - imports = import_matches[:10] # Limit - if not names: # Pure import block - chunk_type = ChunkType.IMPORTS - - # Check for documentation files - if language in ('markdown', 'rst', 'text'): - chunk_type = ChunkType.DOCUMENTATION - - # Check for config files - if language in ('json', 'yaml', 'yml', 'toml', 'xml', 'ini'): - chunk_type = ChunkType.CONFIG - - # Extract docstring if present - docstring = self._extract_docstring(chunk, language) - - # Extract function signature - signature = self._extract_signature(chunk, language) - - # Update metadata - metadata['chunk_type'] = chunk_type.value - - if names: - metadata['semantic_names'] = names[:5] # Limit to 5 names - metadata['primary_name'] = names[0] - - if imports: - metadata['imports'] = imports - - if docstring: - metadata['docstring'] = docstring[:500] # Limit size - - if signature: - metadata['signature'] = signature - - return metadata - - def _extract_docstring(self, chunk: str, language: str) -> Optional[str]: - """Extract docstring from code chunk""" - if language == 'python': - # Python docstrings - match = re.search(r'"""([\s\S]*?)"""|\'\'\'([\s\S]*?)\'\'\'', chunk) - if match: - return (match.group(1) or match.group(2)).strip() - - elif language in ('javascript', 'typescript', 'java', 'csharp', 'php', 'go'): - # JSDoc / JavaDoc style - match = re.search(r'/\*\*([\s\S]*?)\*/', chunk) - if match: - # Clean up the comment - doc = match.group(1) - doc = re.sub(r'^\s*\*\s?', '', doc, flags=re.MULTILINE) - return doc.strip() - - return None - - def _extract_signature(self, chunk: str, language: str) -> Optional[str]: - """Extract function/method signature from code chunk""" - lines = chunk.split('\n') - - for line in lines[:10]: # Check first 10 lines - line = line.strip() - - if language == 'python': - if line.startswith(('def ', 'async def ')): - # Get full signature including multi-line params - sig = line - if ')' not in sig: - # Multi-line signature - idx = lines.index(line.strip()) if line.strip() in lines else -1 - if idx >= 0: - for next_line in lines[idx+1:idx+5]: - sig += ' ' + next_line.strip() - if ')' in next_line: - break - return sig.split(':')[0] + ':' - - elif language in ('java', 'csharp', 'kotlin'): - if any(kw in line for kw in ['public ', 'private ', 'protected ', 'internal ']): - if '(' in line and not line.startswith('//'): - return line.split('{')[0].strip() - - elif language in ('javascript', 'typescript'): - if line.startswith(('function ', 'async function ')): - return line.split('{')[0].strip() - if '=>' in line and '(' in line: - return line.split('=>')[0].strip() + ' =>' - - elif language == 'go': - if line.startswith('func '): - return line.split('{')[0].strip() - - elif language == 'rust': - if line.startswith(('fn ', 'pub fn ', 'async fn ', 'pub async fn ')): - return line.split('{')[0].strip() - - return None - - def _fallback_split(self, doc: Document) -> List[TextNode]: - """Fallback splitting for problematic documents""" - text = doc.text - - if not text or not text.strip(): - return [] - - # Use default splitter - chunks = self._default_splitter.split_text(text) - - nodes = [] - for i, chunk in enumerate(chunks): - if not chunk or not chunk.strip(): - continue - - # Truncate if too large - if len(chunk) > 30000: - chunk = chunk[:30000] - - metadata = dict(doc.metadata) - metadata['chunk_index'] = i - metadata['total_chunks'] = len(chunks) - metadata['chunk_type'] = 'fallback' - - chunk_id = self._make_deterministic_id( - metadata.get('namespace', ''), - metadata.get('path', ''), - i - ) - nodes.append(TextNode( - id_=chunk_id, - text=chunk, - metadata=metadata - )) - - return nodes - - @staticmethod - def get_supported_languages() -> List[str]: - """Return list of supported languages""" - return list(LANGUAGE_MAP.keys()) - - @staticmethod - def get_separators_for_language(language: str) -> Optional[List[str]]: - """Get the separators used for a specific language""" - lang_enum = LANGUAGE_MAP.get(language.lower()) - if lang_enum: - return RecursiveCharacterTextSplitter.get_separators_for_language(lang_enum) - return None diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/__init__.py new file mode 100644 index 00000000..74a91cf0 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/__init__.py @@ -0,0 +1,53 @@ +""" +AST-based code splitter module using Tree-sitter. + +Provides semantic code chunking with: +- Tree-sitter query-based extraction (.scm files) +- Fallback to manual AST traversal +- RecursiveCharacterTextSplitter for oversized chunks +- Rich metadata extraction for RAG +""" + +from .splitter import ASTCodeSplitter, ASTChunk, generate_deterministic_id, compute_file_hash +from .languages import ( + get_language_from_path, + get_treesitter_name, + is_ast_supported, + get_supported_languages, + EXTENSION_TO_LANGUAGE, + AST_SUPPORTED_LANGUAGES, + LANGUAGE_TO_TREESITTER, +) +from .metadata import ContentType, ChunkMetadata, MetadataExtractor +from .tree_parser import TreeSitterParser, get_parser +from .query_runner import QueryRunner, QueryMatch, CapturedNode, get_query_runner + +__all__ = [ + # Main splitter + "ASTCodeSplitter", + "ASTChunk", + "generate_deterministic_id", + "compute_file_hash", + + # Languages + "get_language_from_path", + "get_treesitter_name", + "is_ast_supported", + "get_supported_languages", + "EXTENSION_TO_LANGUAGE", + "AST_SUPPORTED_LANGUAGES", + "LANGUAGE_TO_TREESITTER", + + # Metadata + "ContentType", + "ChunkMetadata", + "MetadataExtractor", + + # Tree-sitter + "TreeSitterParser", + "get_parser", + "QueryRunner", + "QueryMatch", + "CapturedNode", + "get_query_runner", +] diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/languages.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/languages.py new file mode 100644 index 00000000..f4758ff1 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/languages.py @@ -0,0 +1,139 @@ +""" +Language detection and mapping for AST-based code splitting. + +Maps file extensions to tree-sitter language names and LangChain Language enum. +""" + +from pathlib import Path +from typing import Dict, Optional, Set +from langchain_text_splitters import Language + + +# Map file extensions to LangChain Language enum (for RecursiveCharacterTextSplitter fallback) +EXTENSION_TO_LANGUAGE: Dict[str, Language] = { + # Python + '.py': Language.PYTHON, + '.pyw': Language.PYTHON, + '.pyi': Language.PYTHON, + + # Java/JVM + '.java': Language.JAVA, + '.kt': Language.KOTLIN, + '.kts': Language.KOTLIN, + '.scala': Language.SCALA, + + # JavaScript/TypeScript + '.js': Language.JS, + '.jsx': Language.JS, + '.mjs': Language.JS, + '.cjs': Language.JS, + '.ts': Language.TS, + '.tsx': Language.TS, + + # Systems languages + '.go': Language.GO, + '.rs': Language.RUST, + '.c': Language.C, + '.h': Language.C, + '.cpp': Language.CPP, + '.cc': Language.CPP, + '.cxx': Language.CPP, + '.hpp': Language.CPP, + '.hxx': Language.CPP, + '.cs': Language.CSHARP, + + # Web/Scripting + '.php': Language.PHP, + '.phtml': Language.PHP, + '.php3': Language.PHP, + '.php4': Language.PHP, + '.php5': Language.PHP, + '.phps': Language.PHP, + '.inc': Language.PHP, + '.rb': Language.RUBY, + '.erb': Language.RUBY, + '.lua': Language.LUA, + '.pl': Language.PERL, + '.pm': Language.PERL, + '.swift': Language.SWIFT, + + # Markup/Config + '.md': Language.MARKDOWN, + '.markdown': Language.MARKDOWN, + '.html': Language.HTML, + '.htm': Language.HTML, + '.rst': Language.RST, + '.tex': Language.LATEX, + '.proto': Language.PROTO, + '.sol': Language.SOL, + '.hs': Language.HASKELL, + '.cob': Language.COBOL, + '.cbl': Language.COBOL, + '.xml': Language.HTML, +} + +# Languages that support full AST parsing via tree-sitter +AST_SUPPORTED_LANGUAGES: Set[Language] = { + Language.PYTHON, Language.JAVA, Language.KOTLIN, Language.JS, Language.TS, + Language.GO, Language.RUST, Language.C, Language.CPP, Language.CSHARP, + Language.PHP, Language.RUBY, Language.SCALA, Language.LUA, Language.PERL, + Language.SWIFT, Language.HASKELL, Language.COBOL +} + +# Map LangChain Language enum to tree-sitter language name +LANGUAGE_TO_TREESITTER: Dict[Language, str] = { + Language.PYTHON: 'python', + Language.JAVA: 'java', + Language.KOTLIN: 'kotlin', + Language.JS: 'javascript', + Language.TS: 'typescript', + Language.GO: 'go', + Language.RUST: 'rust', + Language.C: 'c', + Language.CPP: 'cpp', + Language.CSHARP: 'c_sharp', + Language.PHP: 'php', + Language.RUBY: 'ruby', + Language.SCALA: 'scala', + Language.LUA: 'lua', + Language.PERL: 'perl', + Language.SWIFT: 'swift', + Language.HASKELL: 'haskell', +} + +# Map tree-sitter language name to module info: (module_name, function_name) +TREESITTER_MODULES: Dict[str, tuple] = { + 'python': ('tree_sitter_python', 'language'), + 'java': ('tree_sitter_java', 'language'), + 'javascript': ('tree_sitter_javascript', 'language'), + 'typescript': ('tree_sitter_typescript', 'language_typescript'), + 'go': ('tree_sitter_go', 'language'), + 'rust': ('tree_sitter_rust', 'language'), + 'c': ('tree_sitter_c', 'language'), + 'cpp': ('tree_sitter_cpp', 'language'), + 'c_sharp': ('tree_sitter_c_sharp', 'language'), + 'ruby': ('tree_sitter_ruby', 'language'), + 'php': ('tree_sitter_php', 'language_php'), +} + + +def get_language_from_path(path: str) -> Optional[Language]: + """Determine LangChain Language enum from file path.""" + ext = Path(path).suffix.lower() + return EXTENSION_TO_LANGUAGE.get(ext) + + +def get_treesitter_name(language: Language) -> Optional[str]: + """Get tree-sitter language name from LangChain Language enum.""" + return LANGUAGE_TO_TREESITTER.get(language) + + +def is_ast_supported(path: str) -> bool: + """Check if AST parsing is supported for a file.""" + language = get_language_from_path(path) + return language is not None and language in AST_SUPPORTED_LANGUAGES + + +def get_supported_languages() -> list: + """Return list of languages with AST support.""" + return list(LANGUAGE_TO_TREESITTER.values()) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/metadata.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/metadata.py new file mode 100644 index 00000000..bcb8559e --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/metadata.py @@ -0,0 +1,343 @@ +""" +Metadata extraction from AST chunks. + +Extracts semantic metadata like docstrings, signatures, inheritance info +from parsed code chunks for improved RAG retrieval. +""" + +import re +import logging +from typing import Dict, Any, List, Optional +from dataclasses import dataclass, field +from enum import Enum + +logger = logging.getLogger(__name__) + + +class ContentType(Enum): + """Content type as determined by AST parsing.""" + FUNCTIONS_CLASSES = "functions_classes" + SIMPLIFIED_CODE = "simplified_code" + FALLBACK = "fallback" + OVERSIZED_SPLIT = "oversized_split" + + +@dataclass +class ChunkMetadata: + """Structured metadata for a code chunk.""" + content_type: ContentType + language: str + path: str + semantic_names: List[str] = field(default_factory=list) + parent_context: List[str] = field(default_factory=list) + docstring: Optional[str] = None + signature: Optional[str] = None + start_line: int = 0 + end_line: int = 0 + node_type: Optional[str] = None + # Class-level metadata + extends: List[str] = field(default_factory=list) + implements: List[str] = field(default_factory=list) + # File-level metadata + imports: List[str] = field(default_factory=list) + namespace: Optional[str] = None + + +class MetadataExtractor: + """ + Extract semantic metadata from code chunks. + + Uses both AST-derived information and regex fallbacks for + comprehensive metadata extraction. + """ + + # Comment prefixes by language + COMMENT_PREFIX: Dict[str, str] = { + 'python': '#', + 'javascript': '//', + 'typescript': '//', + 'java': '//', + 'kotlin': '//', + 'go': '//', + 'rust': '//', + 'c': '//', + 'cpp': '//', + 'c_sharp': '//', + 'php': '//', + 'ruby': '#', + 'lua': '--', + 'perl': '#', + 'scala': '//', + } + + def extract_docstring(self, content: str, language: str) -> Optional[str]: + """Extract docstring from code chunk.""" + if language == 'python': + match = re.search(r'"""([\s\S]*?)"""|\'\'\'([\s\S]*?)\'\'\'', content) + if match: + return (match.group(1) or match.group(2)).strip() + + elif language in ('javascript', 'typescript', 'java', 'kotlin', + 'c_sharp', 'php', 'go', 'scala', 'c', 'cpp'): + # JSDoc / JavaDoc style + match = re.search(r'/\*\*([\s\S]*?)\*/', content) + if match: + doc = match.group(1) + doc = re.sub(r'^\s*\*\s?', '', doc, flags=re.MULTILINE) + return doc.strip() + + elif language == 'rust': + # Rust doc comments + lines = [] + for line in content.split('\n'): + stripped = line.strip() + if stripped.startswith('///'): + lines.append(stripped[3:].strip()) + elif stripped.startswith('//!'): + lines.append(stripped[3:].strip()) + elif lines: + break + if lines: + return '\n'.join(lines) + + return None + + def extract_signature(self, content: str, language: str) -> Optional[str]: + """Extract function/method signature from code chunk.""" + lines = content.split('\n') + + for line in lines[:15]: + line = line.strip() + + if language == 'python': + if line.startswith(('def ', 'async def ', 'class ')): + sig = line + if line.startswith('class ') and ':' in line: + return line.split(':')[0] + ':' + if ')' not in sig and ':' not in sig: + idx = next((i for i, l in enumerate(lines) if l.strip() == line), -1) + if idx >= 0: + for next_line in lines[idx+1:idx+5]: + sig += ' ' + next_line.strip() + if ')' in next_line: + break + if ':' in sig: + return sig.split(':')[0] + ':' + return sig + + elif language in ('java', 'kotlin', 'c_sharp'): + if any(kw in line for kw in ['public ', 'private ', 'protected ', 'internal ', 'fun ']): + if '(' in line and not line.startswith('//'): + return line.split('{')[0].strip() + + elif language in ('javascript', 'typescript'): + if line.startswith(('function ', 'async function ', 'class ')): + return line.split('{')[0].strip() + if '=>' in line and '(' in line: + return line.split('=>')[0].strip() + ' =>' + + elif language == 'go': + if line.startswith('func ') or line.startswith('type '): + return line.split('{')[0].strip() + + elif language == 'rust': + if line.startswith(('fn ', 'pub fn ', 'async fn ', 'pub async fn ', + 'impl ', 'struct ', 'trait ', 'enum ')): + return line.split('{')[0].strip() + + elif language == 'php': + if 'function ' in line and '(' in line: + return line.split('{')[0].strip() + if line.startswith('class ') or line.startswith('interface '): + return line.split('{')[0].strip() + + return None + + def extract_names_from_content(self, content: str, language: str) -> List[str]: + """Extract semantic names (function/class names) using regex patterns.""" + patterns = self._get_name_patterns(language) + names = [] + + for pattern in patterns: + matches = pattern.findall(content) + names.extend(matches) + + # Deduplicate while preserving order + seen = set() + unique_names = [] + for name in names: + if name not in seen: + seen.add(name) + unique_names.append(name) + + return unique_names[:10] # Limit to 10 names + + def _get_name_patterns(self, language: str) -> List[re.Pattern]: + """Get regex patterns for extracting names by language. + + Patterns allow optional leading whitespace (^\s*) to match indented code + like methods inside classes. + """ + patterns = { + 'python': [ + re.compile(r'^\s*class\s+(\w+)', re.MULTILINE), + re.compile(r'^\s*(?:async\s+)?def\s+(\w+)\s*\(', re.MULTILINE), + ], + 'java': [ + re.compile(r'(?:public\s+|private\s+|protected\s+)?(?:abstract\s+|final\s+)?class\s+(\w+)', re.MULTILINE), + re.compile(r'(?:public\s+)?interface\s+(\w+)', re.MULTILINE), + re.compile(r'(?:public|private|protected)\s+(?:static\s+)?[\w<>,\s]+\s+(\w+)\s*\(', re.MULTILINE), + ], + 'javascript': [ + re.compile(r'(?:export\s+)?(?:default\s+)?class\s+(\w+)', re.MULTILINE), + re.compile(r'(?:export\s+)?(?:async\s+)?function\s*\*?\s*(\w+)\s*\(', re.MULTILINE), + re.compile(r'(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\([^)]*\)\s*=>', re.MULTILINE), + ], + 'typescript': [ + re.compile(r'(?:export\s+)?(?:default\s+)?class\s+(\w+)', re.MULTILINE), + re.compile(r'(?:export\s+)?interface\s+(\w+)', re.MULTILINE), + re.compile(r'(?:export\s+)?(?:async\s+)?function\s*\*?\s*(\w+)\s*\(', re.MULTILINE), + re.compile(r'(?:export\s+)?type\s+(\w+)', re.MULTILINE), + ], + 'go': [ + re.compile(r'^\s*func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(', re.MULTILINE), + re.compile(r'^\s*type\s+(\w+)\s+(?:struct|interface)\s*\{', re.MULTILINE), + ], + 'rust': [ + re.compile(r'^\s*(?:pub\s+)?(?:async\s+)?fn\s+(\w+)', re.MULTILINE), + re.compile(r'^\s*(?:pub\s+)?struct\s+(\w+)', re.MULTILINE), + re.compile(r'^\s*(?:pub\s+)?trait\s+(\w+)', re.MULTILINE), + re.compile(r'^\s*(?:pub\s+)?enum\s+(\w+)', re.MULTILINE), + ], + 'php': [ + re.compile(r'(?:abstract\s+|final\s+)?class\s+(\w+)', re.MULTILINE), + re.compile(r'interface\s+(\w+)', re.MULTILINE), + re.compile(r'(?:public|private|protected|static|\s)*function\s+(\w+)\s*\(', re.MULTILINE), + ], + 'c_sharp': [ + re.compile(r'(?:public\s+|private\s+|internal\s+)?(?:abstract\s+|sealed\s+)?class\s+(\w+)', re.MULTILINE), + re.compile(r'(?:public\s+)?interface\s+(\w+)', re.MULTILINE), + re.compile(r'(?:public|private|protected|internal)\s+(?:static\s+)?[\w<>,\s]+\s+(\w+)\s*\(', re.MULTILINE), + ], + } + return patterns.get(language, []) + + def extract_inheritance(self, content: str, language: str) -> Dict[str, List[str]]: + """Extract inheritance information (extends, implements).""" + result = {'extends': [], 'implements': [], 'imports': []} + + patterns = self._get_inheritance_patterns(language) + + if 'extends' in patterns: + match = patterns['extends'].search(content) + if match: + extends = match.group(1).strip() + result['extends'] = [e.strip() for e in extends.split(',') if e.strip()] + + if 'implements' in patterns: + match = patterns['implements'].search(content) + if match: + implements = match.group(1).strip() + result['implements'] = [i.strip() for i in implements.split(',') if i.strip()] + + for key in ('import', 'use', 'using', 'require'): + if key in patterns: + matches = patterns[key].findall(content) + for m in matches: + if isinstance(m, tuple): + result['imports'].extend([x.strip() for x in m if x and x.strip()]) + else: + result['imports'].append(m.strip()) + + # Limit imports + result['imports'] = result['imports'][:20] + + return result + + def _get_inheritance_patterns(self, language: str) -> Dict[str, re.Pattern]: + """Get regex patterns for inheritance extraction.""" + patterns = { + 'python': { + 'extends': re.compile(r'class\s+\w+\s*\(\s*([\w.,\s]+)\s*\)\s*:', re.MULTILINE), + 'import': re.compile(r'^(?:from\s+([\w.]+)\s+)?import\s+([\w.,\s*]+)', re.MULTILINE), + }, + 'java': { + 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), + 'implements': re.compile(r'class\s+\w+(?:\s+extends\s+[\w.]+)?\s+implements\s+([\w.,\s]+)', re.MULTILINE), + 'import': re.compile(r'^import\s+([\w.]+(?:\.\*)?);', re.MULTILINE), + }, + 'typescript': { + 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), + 'implements': re.compile(r'class\s+\w+(?:\s+extends\s+[\w.]+)?\s+implements\s+([\w.,\s]+)', re.MULTILINE), + 'import': re.compile(r'^import\s+(?:[\w{},\s*]+\s+from\s+)?["\']([^"\']+)["\'];?', re.MULTILINE), + }, + 'javascript': { + 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w.]+)', re.MULTILINE), + 'import': re.compile(r'^import\s+(?:[\w{},\s*]+\s+from\s+)?["\']([^"\']+)["\'];?', re.MULTILINE), + 'require': re.compile(r'require\s*\(\s*["\']([^"\']+)["\']\s*\)', re.MULTILINE), + }, + 'php': { + 'extends': re.compile(r'class\s+\w+\s+extends\s+([\w\\]+)', re.MULTILINE), + 'implements': re.compile(r'class\s+\w+(?:\s+extends\s+[\w\\]+)?\s+implements\s+([\w\\,\s]+)', re.MULTILINE), + 'use': re.compile(r'^use\s+([\w\\]+)(?:\s+as\s+\w+)?;', re.MULTILINE), + }, + 'c_sharp': { + 'extends': re.compile(r'class\s+\w+\s*:\s*([\w.]+)', re.MULTILINE), + 'using': re.compile(r'^using\s+([\w.]+);', re.MULTILINE), + }, + 'go': { + 'import': re.compile(r'^import\s+(?:\(\s*)?"([^"]+)"', re.MULTILINE), + }, + 'rust': { + 'use': re.compile(r'^use\s+([\w:]+(?:::\{[^}]+\})?);', re.MULTILINE), + }, + } + return patterns.get(language, {}) + + def get_comment_prefix(self, language: str) -> str: + """Get comment prefix for a language.""" + return self.COMMENT_PREFIX.get(language, '//') + + def build_metadata_dict( + self, + chunk_metadata: ChunkMetadata, + base_metadata: Dict[str, Any] + ) -> Dict[str, Any]: + """Build final metadata dictionary from ChunkMetadata.""" + metadata = dict(base_metadata) + + metadata['content_type'] = chunk_metadata.content_type.value + metadata['node_type'] = chunk_metadata.node_type + metadata['start_line'] = chunk_metadata.start_line + metadata['end_line'] = chunk_metadata.end_line + + if chunk_metadata.parent_context: + metadata['parent_context'] = chunk_metadata.parent_context + metadata['parent_class'] = chunk_metadata.parent_context[-1] + full_path_parts = chunk_metadata.parent_context + chunk_metadata.semantic_names[:1] + metadata['full_path'] = '.'.join(full_path_parts) + + if chunk_metadata.semantic_names: + metadata['semantic_names'] = chunk_metadata.semantic_names + metadata['primary_name'] = chunk_metadata.semantic_names[0] + + if chunk_metadata.docstring: + metadata['docstring'] = chunk_metadata.docstring[:500] + + if chunk_metadata.signature: + metadata['signature'] = chunk_metadata.signature + + if chunk_metadata.extends: + metadata['extends'] = chunk_metadata.extends + metadata['parent_types'] = chunk_metadata.extends + + if chunk_metadata.implements: + metadata['implements'] = chunk_metadata.implements + + if chunk_metadata.imports: + metadata['imports'] = chunk_metadata.imports + + if chunk_metadata.namespace: + metadata['namespace'] = chunk_metadata.namespace + + return metadata diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/c_sharp.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/c_sharp.scm new file mode 100644 index 00000000..03f4775a --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/c_sharp.scm @@ -0,0 +1,63 @@ +; C# tree-sitter queries for AST-based code splitting + +; Using directives +(using_directive) @using + +; Namespace declarations (simple names) +(namespace_declaration + name: (identifier) @name) @definition.namespace + +; Namespace declarations (qualified names like My.App) +(namespace_declaration + name: (qualified_name) @name) @definition.namespace + +(file_scoped_namespace_declaration + name: (identifier) @name) @definition.namespace + +(file_scoped_namespace_declaration + name: (qualified_name) @name) @definition.namespace + +; Class declarations +(class_declaration + name: (identifier) @name) @definition.class + +; Struct declarations +(struct_declaration + name: (identifier) @name) @definition.struct + +; Interface declarations +(interface_declaration + name: (identifier) @name) @definition.interface + +; Enum declarations +(enum_declaration + name: (identifier) @name) @definition.enum + +; Record declarations +(record_declaration + name: (identifier) @name) @definition.record + +; Delegate declarations +(delegate_declaration + name: (identifier) @name) @definition.delegate + +; Method declarations +(method_declaration + name: (identifier) @name) @definition.method + +; Constructor declarations +(constructor_declaration + name: (identifier) @name) @definition.constructor + +; Property declarations +(property_declaration + name: (identifier) @name) @definition.property + +; Field declarations +(field_declaration) @definition.field + +; Event declarations +(event_declaration) @definition.event + +; Attributes +(attribute) @attribute diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/go.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/go.scm new file mode 100644 index 00000000..87bf86ff --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/go.scm @@ -0,0 +1,26 @@ +; Go tree-sitter queries for AST-based code splitting + +; Package clause +(package_clause) @package + +; Import declarations +(import_declaration) @import + +; Type declarations (struct, interface, type alias) +(type_declaration + (type_spec + name: (type_identifier) @name)) @definition.type + +; Function declarations +(function_declaration + name: (identifier) @name) @definition.function + +; Method declarations +(method_declaration + name: (field_identifier) @name) @definition.method + +; Variable declarations +(var_declaration) @definition.variable + +; Constant declarations +(const_declaration) @definition.const diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/java.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/java.scm new file mode 100644 index 00000000..aebcac8c --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/java.scm @@ -0,0 +1,52 @@ +; Java tree-sitter queries for AST-based code splitting + +; Package declaration +(package_declaration) @package + +; Import declarations +(import_declaration) @import + +; Class declarations +(class_declaration + name: (identifier) @name) @definition.class + +; Interface declarations +(interface_declaration + name: (identifier) @name) @definition.interface + +; Enum declarations +(enum_declaration + name: (identifier) @name) @definition.enum + +; Record declarations (Java 14+) +(record_declaration + name: (identifier) @name) @definition.record + +; Annotation type declarations +(annotation_type_declaration + name: (identifier) @name) @definition.annotation + +; Method declarations +(method_declaration + name: (identifier) @name) @definition.method + +; Constructor declarations +(constructor_declaration + name: (identifier) @name) @definition.constructor + +; Field declarations +(field_declaration) @definition.field + +; Annotations (for metadata) - simple names +(marker_annotation + name: (identifier) @name) @annotation + +(annotation + name: (identifier) @name) @annotation + +; Annotations with fully qualified names (e.g., @org.junit.Test) +(marker_annotation + name: (scoped_identifier) @name) @annotation + +(annotation + name: (scoped_identifier) @name) @annotation diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/javascript.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/javascript.scm new file mode 100644 index 00000000..dd8c3561 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/javascript.scm @@ -0,0 +1,42 @@ +; JavaScript/TypeScript tree-sitter queries for AST-based code splitting + +; Import statements +(import_statement) @import + +; Export statements +(export_statement) @export + +; Class declarations +(class_declaration + name: (identifier) @name) @definition.class + +; Function declarations +(function_declaration + name: (identifier) @name) @definition.function + +; Arrow functions assigned to variables +(lexical_declaration + (variable_declarator + name: (identifier) @name + value: (arrow_function))) @definition.function + +; Generator functions +(generator_function_declaration + name: (identifier) @name) @definition.function + +; Method definitions (inside class body) +(method_definition + name: (property_identifier) @name) @definition.method + +; Variable declarations (module-level) +(lexical_declaration) @definition.variable + +(variable_declaration) @definition.variable + +; Interface declarations (TypeScript) +(interface_declaration + name: (type_identifier) @name) @definition.interface + +; Type alias declarations (TypeScript) +(type_alias_declaration + name: (type_identifier) @name) @definition.type diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/php.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/php.scm new file mode 100644 index 00000000..52c29810 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/php.scm @@ -0,0 +1,40 @@ +; PHP tree-sitter queries for AST-based code splitting + +; Namespace definition +(namespace_definition) @namespace + +; Use statements +(namespace_use_declaration) @use + +; Class declarations +(class_declaration + name: (name) @name) @definition.class + +; Interface declarations +(interface_declaration + name: (name) @name) @definition.interface + +; Trait declarations +(trait_declaration + name: (name) @name) @definition.trait + +; Enum declarations (PHP 8.1+) +(enum_declaration + name: (name) @name) @definition.enum + +; Function definitions +(function_definition + name: (name) @name) @definition.function + +; Method declarations +(method_declaration + name: (name) @name) @definition.method + +; Property declarations +(property_declaration) @definition.property + +; Const declarations +(const_declaration) @definition.const + +; Attributes (PHP 8.0+) +(attribute) @attribute diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/python.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/python.scm new file mode 100644 index 00000000..3429a53c --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/python.scm @@ -0,0 +1,29 @@ +; Python tree-sitter queries for AST-based code splitting + +; Import statements +(import_statement) @import + +(import_from_statement) @import + +; Class definitions +(class_definition + name: (identifier) @name) @definition.class + +; Function definitions +(function_definition + name: (identifier) @name) @definition.function + +; Decorated definitions (class or function with decorators) +(decorated_definition) @definition.decorated + +; Decorators +(decorator) @decorator + +; Assignment statements (module-level constants) +(assignment + left: (identifier) @name) @definition.assignment + +; Type alias (Python 3.12+) +; tree-sitter grammar uses field "left" with type "type" for the alias name +(type_alias_statement + left: (type) @name) @definition.type_alias diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/rust.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/rust.scm new file mode 100644 index 00000000..045150be --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/rust.scm @@ -0,0 +1,52 @@ +; Rust tree-sitter queries for AST-based code splitting + +; Use declarations +(use_declaration) @use + +; Module declarations +(mod_item + name: (identifier) @name) @definition.module + +; Struct definitions +(struct_item + name: (type_identifier) @name) @definition.struct + +; Enum definitions +(enum_item + name: (type_identifier) @name) @definition.enum + +; Trait definitions +(trait_item + name: (type_identifier) @name) @definition.trait + +; Implementation blocks (simple type impl) +(impl_item + type: (type_identifier) @name) @definition.impl + +; Implementation blocks for trait (impl Trait for Type) +(impl_item + trait: (type_identifier) @trait_name + type: (type_identifier) @name) @definition.impl + +; Function definitions +(function_item + name: (identifier) @name) @definition.function + +; Type alias +(type_item + name: (type_identifier) @name) @definition.type + +; Constant definitions +(const_item + name: (identifier) @name) @definition.const + +; Static definitions +(static_item + name: (identifier) @name) @definition.static + +; Macro definitions +(macro_definition + name: (identifier) @name) @definition.macro + +; Attributes +(attribute_item) @attribute diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/typescript.scm b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/typescript.scm new file mode 100644 index 00000000..b65f43cc --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/queries/typescript.scm @@ -0,0 +1,61 @@ +; TypeScript tree-sitter queries for AST-based code splitting +; Uses same patterns as JavaScript plus TypeScript-specific nodes + +; Import statements +(import_statement) @import + +; Export statements +(export_statement) @export + +; Class declarations +(class_declaration + name: (type_identifier) @name) @definition.class + +; Abstract class declarations +(abstract_class_declaration + name: (type_identifier) @name) @definition.class + +; Function declarations +(function_declaration + name: (identifier) @name) @definition.function + +; Arrow functions assigned to variables +(lexical_declaration + (variable_declarator + name: (identifier) @name + value: (arrow_function))) @definition.function + +; Method definitions (inside class body) +(method_definition + name: (property_identifier) @name) @definition.method + +; Interface declarations +(interface_declaration + name: (type_identifier) @name) @definition.interface + +; Type alias declarations +(type_alias_declaration + name: (type_identifier) @name) @definition.type + +; Enum declarations +(enum_declaration + name: (identifier) @name) @definition.enum + +; Module declarations (ambient/string-named modules) +(module + name: (identifier) @name) @definition.module + +; TypeScript namespace declarations (internal modules) +(internal_module + name: (identifier) @name) @definition.namespace + +; Variable declarations at module level (excluding arrow functions captured above) +; Only capture if NOT an arrow function assignment +(lexical_declaration + (variable_declarator + name: (identifier) @name + value: (_) @value) @declarator + (#not-match? @value "^arrow_function")) @definition.variable + +; Ambient declarations +(ambient_declaration) @definition.ambient diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/query_runner.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/query_runner.py new file mode 100644 index 00000000..7ef2bda2 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/query_runner.py @@ -0,0 +1,350 @@ +""" +Tree-sitter query runner using custom query files with built-in fallback. + +Prefers custom .scm query files for rich metadata extraction (extends, implements, imports), +falling back to built-in TAGS_QUERY only when custom query is unavailable. +""" + +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Any, Optional + +from .tree_parser import get_parser +from .languages import TREESITTER_MODULES + +logger = logging.getLogger(__name__) + +# Directory containing custom .scm query files +QUERIES_DIR = Path(__file__).parent / "queries" + +# Languages that have built-in TAGS_QUERY (used as fallback only) +LANGUAGES_WITH_BUILTIN_TAGS = {'python', 'java', 'javascript', 'go', 'rust', 'php'} + +# Languages with custom .scm files for rich metadata (extends, implements, imports) +LANGUAGES_WITH_CUSTOM_QUERY = { + 'python', 'java', 'javascript', 'typescript', 'c_sharp', 'go', 'rust', 'php' +} + + +@dataclass +class CapturedNode: + """Represents a captured AST node from a query.""" + name: str # Capture name (e.g., 'function.name', 'class.body') + text: str # Node text content + start_byte: int + end_byte: int + start_point: tuple # (row, column) + end_point: tuple + node_type: str # Tree-sitter node type + + @property + def start_line(self) -> int: + return self.start_point[0] + 1 # Convert to 1-based + + @property + def end_line(self) -> int: + return self.end_point[0] + 1 + + +@dataclass +class QueryMatch: + """A complete match from a query pattern.""" + pattern_name: str # e.g., 'function', 'class', 'import' + captures: Dict[str, CapturedNode] = field(default_factory=dict) + + def get(self, capture_name: str) -> Optional[CapturedNode]: + """Get a captured node by name.""" + return self.captures.get(capture_name) + + @property + def full_text(self) -> Optional[str]: + """Get the full text of the main capture (pattern_name without suffix).""" + main_capture = self.captures.get(self.pattern_name) + return main_capture.text if main_capture else None + + +class QueryRunner: + """ + Executes tree-sitter queries using custom .scm files with built-in fallback. + + Strategy: + 1. Prefer custom .scm files for rich metadata (extends, implements, imports, decorators) + 2. Fall back to built-in TAGS_QUERY only when no custom query exists + + Custom queries capture: @class.extends, @class.implements, @import, @decorator, + @method.visibility, @function.return_type, etc. + + Built-in TAGS_QUERY only captures: @definition.function, @definition.class, @name, @doc + """ + + def __init__(self): + self._query_cache: Dict[str, Any] = {} # lang -> compiled query + self._scm_cache: Dict[str, str] = {} # lang -> raw scm string + self._parser = get_parser() + + def _get_builtin_tags_query(self, lang_name: str) -> Optional[str]: + """Get built-in TAGS_QUERY from language package if available.""" + if lang_name not in LANGUAGES_WITH_BUILTIN_TAGS: + return None + + lang_info = TREESITTER_MODULES.get(lang_name) + if not lang_info: + return None + + module_name = lang_info[0] + try: + import importlib + lang_module = importlib.import_module(module_name) + tags_query = getattr(lang_module, 'TAGS_QUERY', None) + if tags_query: + logger.debug(f"Using built-in TAGS_QUERY for {lang_name}") + return tags_query + except (ImportError, AttributeError) as e: + logger.debug(f"Could not load built-in query for {lang_name}: {e}") + + return None + + def _load_custom_query_file(self, lang_name: str) -> Optional[str]: + """Load custom .scm query file for languages without built-in queries.""" + if lang_name in self._scm_cache: + return self._scm_cache[lang_name] + + query_file = QUERIES_DIR / f"{lang_name}.scm" + + if not query_file.exists(): + logger.debug(f"No custom query file for {lang_name}") + return None + + try: + scm_content = query_file.read_text(encoding='utf-8') + self._scm_cache[lang_name] = scm_content + logger.debug(f"Loaded custom query file for {lang_name}") + return scm_content + except Exception as e: + logger.warning(f"Failed to load query file {query_file}: {e}") + return None + + def _try_compile_query(self, lang_name: str, scm_content: str, language: Any) -> Optional[Any]: + """Try to compile a query string, returning None on failure.""" + try: + from tree_sitter import Query + return Query(language, scm_content) + except Exception as e: + logger.debug(f"Query compilation failed for {lang_name}: {e}") + return None + + def _get_compiled_query(self, lang_name: str) -> Optional[Any]: + """Get or compile the query for a language with fallback.""" + if lang_name in self._query_cache: + return self._query_cache[lang_name] + + language = self._parser.get_language(lang_name) + if not language: + return None + + # Try custom .scm first + custom_scm = self._load_custom_query_file(lang_name) + if custom_scm: + query = self._try_compile_query(lang_name, custom_scm, language) + if query: + logger.debug(f"Using custom query for {lang_name}") + self._query_cache[lang_name] = query + return query + else: + logger.debug(f"Custom query failed for {lang_name}, trying built-in") + + # Fallback to built-in TAGS_QUERY + builtin_scm = self._get_builtin_tags_query(lang_name) + if builtin_scm: + query = self._try_compile_query(lang_name, builtin_scm, language) + if query: + logger.debug(f"Using built-in TAGS_QUERY for {lang_name}") + self._query_cache[lang_name] = query + return query + + logger.debug(f"No working query available for {lang_name}") + return None + + def run_query( + self, + source_code: str, + lang_name: str, + tree: Optional[Any] = None + ) -> List[QueryMatch]: + """ + Run the query for a language and return all matches. + + Args: + source_code: Source code string + lang_name: Tree-sitter language name + tree: Optional pre-parsed tree (will parse if not provided) + + Returns: + List of QueryMatch objects with captured nodes + """ + query = self._get_compiled_query(lang_name) + if not query: + return [] + + if tree is None: + tree = self._parser.parse(source_code, lang_name) + if tree is None: + return [] + + source_bytes = source_code.encode('utf-8') + + try: + # Use QueryCursor.matches() for pattern-grouped results + # Each match is (pattern_id, {capture_name: [nodes]}) + from tree_sitter import QueryCursor + cursor = QueryCursor(query) + raw_matches = list(cursor.matches(tree.root_node)) + except Exception as e: + logger.warning(f"Query execution failed for {lang_name}: {e}") + return [] + + results: List[QueryMatch] = [] + + for pattern_id, captures_dict in raw_matches: + # Determine pattern type from captures + # Built-in: @definition.function, @definition.class, @name + # Custom: @function, @class, @function.name + + pattern_name = None + main_node = None + name_node = None + doc_node = None + + for capture_name, nodes in captures_dict.items(): + if not nodes: + continue + node = nodes[0] # Take first node for each capture + + # Built-in definition captures + if capture_name.startswith('definition.'): + pattern_name = capture_name[len('definition.'):] + main_node = node + # Built-in @name capture (associated with this pattern) + elif capture_name == 'name': + name_node = node + # Built-in @doc capture + elif capture_name == 'doc': + doc_node = node + # Skip reference captures + elif capture_name.startswith('reference.'): + continue + # Custom query captures: @function, @class + elif '.' not in capture_name: + pattern_name = capture_name + main_node = node + + # Skip if no definition pattern found + if not pattern_name or not main_node: + continue + + # Build the QueryMatch + match = QueryMatch(pattern_name=pattern_name) + + # Add main capture + match.captures[pattern_name] = CapturedNode( + name=pattern_name, + text=source_bytes[main_node.start_byte:main_node.end_byte].decode('utf-8', errors='replace'), + start_byte=main_node.start_byte, + end_byte=main_node.end_byte, + start_point=(main_node.start_point.row, main_node.start_point.column), + end_point=(main_node.end_point.row, main_node.end_point.column), + node_type=main_node.type + ) + + # Add name capture if present + if name_node: + match.captures[f'{pattern_name}.name'] = CapturedNode( + name=f'{pattern_name}.name', + text=source_bytes[name_node.start_byte:name_node.end_byte].decode('utf-8', errors='replace'), + start_byte=name_node.start_byte, + end_byte=name_node.end_byte, + start_point=(name_node.start_point.row, name_node.start_point.column), + end_point=(name_node.end_point.row, name_node.end_point.column), + node_type=name_node.type + ) + + # Add doc capture if present + if doc_node: + match.captures[f'{pattern_name}.doc'] = CapturedNode( + name=f'{pattern_name}.doc', + text=source_bytes[doc_node.start_byte:doc_node.end_byte].decode('utf-8', errors='replace'), + start_byte=doc_node.start_byte, + end_byte=doc_node.end_byte, + start_point=(doc_node.start_point.row, doc_node.start_point.column), + end_point=(doc_node.end_point.row, doc_node.end_point.column), + node_type=doc_node.type + ) + + # Process any additional sub-captures from custom queries + for capture_name, nodes in captures_dict.items(): + if '.' in capture_name and not capture_name.startswith(('definition.', 'reference.')): + node = nodes[0] + match.captures[capture_name] = CapturedNode( + name=capture_name, + text=source_bytes[node.start_byte:node.end_byte].decode('utf-8', errors='replace'), + start_byte=node.start_byte, + end_byte=node.end_byte, + start_point=(node.start_point.row, node.start_point.column), + end_point=(node.end_point.row, node.end_point.column), + node_type=node.type + ) + + results.append(match) + + return results + + def get_functions(self, source_code: str, lang_name: str) -> List[QueryMatch]: + """Convenience method to get function/method matches.""" + matches = self.run_query(source_code, lang_name) + return [m for m in matches if m.pattern_name in ('function', 'method')] + + def get_classes(self, source_code: str, lang_name: str) -> List[QueryMatch]: + """Convenience method to get class/struct/interface matches.""" + matches = self.run_query(source_code, lang_name) + return [m for m in matches if m.pattern_name in ('class', 'struct', 'interface', 'trait')] + + def get_imports(self, source_code: str, lang_name: str) -> List[QueryMatch]: + """Convenience method to get import statement matches.""" + matches = self.run_query(source_code, lang_name) + return [m for m in matches if m.pattern_name == 'import'] + + def has_query(self, lang_name: str) -> bool: + """Check if a query is available for this language (custom or built-in).""" + # Check custom file first + query_file = QUERIES_DIR / f"{lang_name}.scm" + if query_file.exists(): + return True + # Check built-in fallback + return lang_name in LANGUAGES_WITH_BUILTIN_TAGS + + def uses_custom_query(self, lang_name: str) -> bool: + """Check if this language uses custom .scm query (rich metadata).""" + query_file = QUERIES_DIR / f"{lang_name}.scm" + return query_file.exists() + + def uses_builtin_query(self, lang_name: str) -> bool: + """Check if this language uses built-in TAGS_QUERY (limited metadata).""" + return lang_name in LANGUAGES_WITH_BUILTIN_TAGS and not self.uses_custom_query(lang_name) + + def clear_cache(self): + """Clear compiled query cache.""" + self._query_cache.clear() + self._scm_cache.clear() + + +# Global singleton +_runner_instance: Optional[QueryRunner] = None + + +def get_query_runner() -> QueryRunner: + """Get the global QueryRunner instance.""" + global _runner_instance + if _runner_instance is None: + _runner_instance = QueryRunner() + return _runner_instance diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/splitter.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/splitter.py new file mode 100644 index 00000000..25c12d8a --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/splitter.py @@ -0,0 +1,1326 @@ +""" +AST-based Code Splitter using Tree-sitter for accurate code parsing. + +This module provides true AST-aware code chunking that: +1. Uses Tree-sitter queries for efficient pattern matching (15+ languages) +2. Splits code into semantic units (classes, functions, methods) +3. Uses RecursiveCharacterTextSplitter for oversized chunks +4. Enriches metadata for better RAG retrieval +5. Maintains parent context ("breadcrumbs") for nested structures +6. Uses deterministic IDs for Qdrant deduplication +""" + +import hashlib +import logging +from typing import List, Dict, Any, Optional, Set +from pathlib import Path +from dataclasses import dataclass, field + +from langchain_text_splitters import RecursiveCharacterTextSplitter, Language +from llama_index.core.schema import Document as LlamaDocument, TextNode + +from .languages import ( + EXTENSION_TO_LANGUAGE, AST_SUPPORTED_LANGUAGES, LANGUAGE_TO_TREESITTER, + get_language_from_path, get_treesitter_name, is_ast_supported +) +from .tree_parser import get_parser +from .query_runner import get_query_runner, QueryMatch +from .metadata import MetadataExtractor, ContentType, ChunkMetadata + +logger = logging.getLogger(__name__) + + +def generate_deterministic_id(path: str, content: str, chunk_index: int = 0) -> str: + """ + Generate a deterministic ID for a chunk based on file path and content. + + This ensures the same code chunk always gets the same ID, preventing + duplicates in Qdrant during re-indexing. + """ + hash_input = f"{path}:{chunk_index}:{content[:500]}" + return hashlib.sha256(hash_input.encode('utf-8')).hexdigest()[:32] + + +def compute_file_hash(content: str) -> str: + """Compute hash of file content for change detection.""" + return hashlib.sha256(content.encode('utf-8')).hexdigest() + + +@dataclass +class ASTChunk: + """Represents a chunk of code from AST parsing with rich metadata.""" + content: str + content_type: ContentType + language: str + path: str + + # Identity + semantic_names: List[str] = field(default_factory=list) + node_type: Optional[str] = None + namespace: Optional[str] = None + + # Location + start_line: int = 0 + end_line: int = 0 + + # Hierarchy & Context + parent_context: List[str] = field(default_factory=list) # Breadcrumb path + + # Documentation + docstring: Optional[str] = None + signature: Optional[str] = None + + # Type relationships + extends: List[str] = field(default_factory=list) + implements: List[str] = field(default_factory=list) + + # Dependencies + imports: List[str] = field(default_factory=list) + + # --- RICH AST FIELDS (extracted from tree-sitter) --- + + # Methods/functions within this chunk (for classes) + methods: List[str] = field(default_factory=list) + + # Properties/fields within this chunk (for classes) + properties: List[str] = field(default_factory=list) + + # Parameters (for functions/methods) + parameters: List[str] = field(default_factory=list) + + # Return type (for functions/methods) + return_type: Optional[str] = None + + # Decorators/annotations + decorators: List[str] = field(default_factory=list) + + # Modifiers (public, private, static, async, abstract, etc.) + modifiers: List[str] = field(default_factory=list) + + # Called functions/methods (dependencies) + calls: List[str] = field(default_factory=list) + + # Referenced types (type annotations, generics) + referenced_types: List[str] = field(default_factory=list) + + # Variables declared in this chunk + variables: List[str] = field(default_factory=list) + + # Constants defined + constants: List[str] = field(default_factory=list) + + # Generic type parameters (e.g., ) + type_parameters: List[str] = field(default_factory=list) + + +class ASTCodeSplitter: + """ + AST-based code splitter using Tree-sitter queries for accurate parsing. + + Features: + - Uses .scm query files for declarative pattern matching + - Splits code into semantic units (classes, functions, methods) + - Falls back to RecursiveCharacterTextSplitter when needed + - Uses deterministic IDs for Qdrant deduplication + - Enriches metadata for improved RAG retrieval + - Prepares embedding-optimized text with semantic context + + Chunk Size Strategy: + - text-embedding-3-small supports ~8191 tokens (~32K chars) + - We use 8000 chars as default to keep semantic units intact + - Only truly massive classes/functions get split + - Splitting loses AST benefits, so we avoid it when possible + + Usage: + splitter = ASTCodeSplitter(max_chunk_size=8000) + nodes = splitter.split_documents(documents) + """ + + # Chunk size considerations: + # - Embedding models (text-embedding-3-small): ~8191 tokens = ~32K chars + # - Most classes/functions: 500-5000 chars + # - Keeping semantic units whole improves retrieval quality + # - Only split when absolutely necessary + DEFAULT_MAX_CHUNK_SIZE = 8000 # ~2000 tokens, fits most semantic units + DEFAULT_MIN_CHUNK_SIZE = 100 + DEFAULT_CHUNK_OVERLAP = 200 + DEFAULT_PARSER_THRESHOLD = 3 # Low threshold - AST benefits even small files + + def __init__( + self, + max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE, + min_chunk_size: int = DEFAULT_MIN_CHUNK_SIZE, + chunk_overlap: int = DEFAULT_CHUNK_OVERLAP, + parser_threshold: int = DEFAULT_PARSER_THRESHOLD, + enrich_embedding_text: bool = True + ): + """ + Initialize AST code splitter. + + Args: + max_chunk_size: Maximum characters per chunk + min_chunk_size: Minimum characters for a valid chunk + chunk_overlap: Overlap between chunks when splitting oversized content + parser_threshold: Minimum lines for AST parsing (3 recommended) + enrich_embedding_text: Whether to prepend semantic context to chunk text + for better embedding quality + """ + self.max_chunk_size = max_chunk_size + self.min_chunk_size = min_chunk_size + self.chunk_overlap = chunk_overlap + self.parser_threshold = parser_threshold + self.enrich_embedding_text = enrich_embedding_text + + # Components + self._parser = get_parser() + self._query_runner = get_query_runner() + self._metadata_extractor = MetadataExtractor() + + # Cache text splitters + self._splitter_cache: Dict[Language, RecursiveCharacterTextSplitter] = {} + + # Default splitter + self._default_splitter = RecursiveCharacterTextSplitter( + chunk_size=max_chunk_size, + chunk_overlap=chunk_overlap, + length_function=len, + ) + + def split_documents(self, documents: List[LlamaDocument]) -> List[TextNode]: + """ + Split LlamaIndex documents using AST-based parsing. + + Args: + documents: List of LlamaIndex Document objects + + Returns: + List of TextNode objects with enriched metadata + """ + all_nodes = [] + + for doc in documents: + path = doc.metadata.get('path', 'unknown') + language = get_language_from_path(path) + + line_count = doc.text.count('\n') + 1 + use_ast = ( + language is not None + and language in AST_SUPPORTED_LANGUAGES + and line_count >= self.parser_threshold + and self._parser.is_available() + ) + + if use_ast: + nodes = self._split_with_ast(doc, language) + else: + nodes = self._split_fallback(doc, language) + + all_nodes.extend(nodes) + logger.debug(f"Split {path} into {len(nodes)} chunks (AST={use_ast})") + + return all_nodes + + def _split_with_ast(self, doc: LlamaDocument, language: Language) -> List[TextNode]: + """Split document using AST parsing with query-based extraction.""" + text = doc.text + path = doc.metadata.get('path', 'unknown') + ts_lang = get_treesitter_name(language) + + if not ts_lang: + return self._split_fallback(doc, language) + + # Try query-based extraction first + chunks = self._extract_with_queries(text, ts_lang, path) + + # If no queries available, fall back to traversal-based extraction + if not chunks: + chunks = self._extract_with_traversal(text, ts_lang, path) + + # Still no chunks? Use fallback + if not chunks: + return self._split_fallback(doc, language) + + return self._process_chunks(chunks, doc, language, path) + + def _extract_with_queries( + self, + text: str, + lang_name: str, + path: str + ) -> List[ASTChunk]: + """Extract chunks using tree-sitter query files with rich metadata.""" + if not self._query_runner.has_query(lang_name): + return [] + + tree = self._parser.parse(text, lang_name) + if not tree: + return [] + + matches = self._query_runner.run_query(text, lang_name, tree) + if not matches: + return [] + + source_bytes = text.encode('utf-8') + chunks = [] + processed_ranges: Set[tuple] = set() + + # Collect file-level metadata from all matches + imports = [] + namespace = None + decorators_map: Dict[int, List[str]] = {} # line -> decorators + + for match in matches: + # Handle imports (multiple capture variations) + if match.pattern_name in ('import', 'use'): + import_path = ( + match.get('import.path') or + match.get('import') or + match.get('use.path') or + match.get('use') + ) + if import_path: + imports.append(import_path.text.strip().strip('"\'')) + continue + + # Handle namespace/package/module + if match.pattern_name in ('namespace', 'package', 'module'): + ns_cap = match.get(f'{match.pattern_name}.name') or match.get(match.pattern_name) + if ns_cap: + namespace = ns_cap.text.strip() + continue + + # Handle standalone decorators/attributes + if match.pattern_name in ('decorator', 'attribute', 'annotation'): + dec_cap = match.get(f'{match.pattern_name}.name') or match.get(match.pattern_name) + if dec_cap: + line = dec_cap.start_line + if line not in decorators_map: + decorators_map[line] = [] + decorators_map[line].append(dec_cap.text.strip()) + continue + + # Handle main constructs: functions, classes, methods, etc. + semantic_patterns = ( + 'function', 'method', 'class', 'interface', 'struct', 'trait', + 'enum', 'impl', 'constructor', 'closure', 'arrow', 'const', + 'var', 'static', 'type', 'record' + ) + if match.pattern_name in semantic_patterns: + main_cap = match.get(match.pattern_name) + if not main_cap: + continue + + range_key = (main_cap.start_byte, main_cap.end_byte) + if range_key in processed_ranges: + continue + processed_ranges.add(range_key) + + # Get name from various capture patterns + name_cap = ( + match.get(f'{match.pattern_name}.name') or + match.get('name') + ) + name = name_cap.text if name_cap else None + + # Get inheritance (extends/implements/embeds/supertrait) + extends = [] + implements = [] + + for ext_capture in ('extends', 'embeds', 'supertrait', 'base_type'): + cap = match.get(f'{match.pattern_name}.{ext_capture}') + if cap: + extends.extend(self._parse_type_list(cap.text)) + + for impl_capture in ('implements', 'trait'): + cap = match.get(f'{match.pattern_name}.{impl_capture}') + if cap: + implements.extend(self._parse_type_list(cap.text)) + + # Get additional metadata from captures + visibility = match.get(f'{match.pattern_name}.visibility') + return_type = match.get(f'{match.pattern_name}.return_type') + params = match.get(f'{match.pattern_name}.params') + modifiers = [] + + for mod in ('static', 'abstract', 'final', 'async', 'readonly', 'const', 'unsafe'): + if match.get(f'{match.pattern_name}.{mod}'): + modifiers.append(mod) + + chunk = ASTChunk( + content=main_cap.text, + content_type=ContentType.FUNCTIONS_CLASSES, + language=lang_name, + path=path, + semantic_names=[name] if name else [], + parent_context=[], + start_line=main_cap.start_line, + end_line=main_cap.end_line, + node_type=match.pattern_name, + extends=extends, + implements=implements, + modifiers=modifiers, + ) + + # Extract docstring and signature + chunk.docstring = self._metadata_extractor.extract_docstring(main_cap.text, lang_name) + chunk.signature = self._metadata_extractor.extract_signature(main_cap.text, lang_name) + + # Extract rich AST details (methods, properties, params, calls, etc.) + self._extract_rich_ast_details(chunk, tree, main_cap, lang_name) + + chunks.append(chunk) + + # Add imports and namespace to all chunks + for chunk in chunks: + chunk.imports = imports[:30] + chunk.namespace = namespace + + # Create simplified code chunk + if chunks: + simplified = self._create_simplified_code(text, chunks, lang_name) + if simplified and len(simplified.strip()) > 50: + chunks.append(ASTChunk( + content=simplified, + content_type=ContentType.SIMPLIFIED_CODE, + language=lang_name, + path=path, + start_line=1, + end_line=text.count('\n') + 1, + node_type='simplified', + imports=imports[:30], + namespace=namespace, + )) + + return chunks + + def _extract_with_traversal( + self, + text: str, + lang_name: str, + path: str + ) -> List[ASTChunk]: + """Fallback: extract chunks using manual AST traversal.""" + tree = self._parser.parse(text, lang_name) + if not tree: + return [] + + source_bytes = text.encode('utf-8') + chunks = [] + processed_ranges: Set[tuple] = set() + + # Node types for semantic chunking + semantic_types = self._get_semantic_node_types(lang_name) + class_types = set(semantic_types.get('class', [])) + function_types = set(semantic_types.get('function', [])) + all_types = class_types | function_types + + def get_node_text(node) -> str: + return source_bytes[node.start_byte:node.end_byte].decode('utf-8', errors='replace') + + def get_node_name(node) -> Optional[str]: + for child in node.children: + if child.type in ('identifier', 'name', 'type_identifier', 'property_identifier'): + return get_node_text(child) + return None + + def traverse(node, parent_context: List[str]): + node_range = (node.start_byte, node.end_byte) + + if node.type in all_types: + if node_range in processed_ranges: + return + + content = get_node_text(node) + start_line = source_bytes[:node.start_byte].count(b'\n') + 1 + end_line = start_line + content.count('\n') + node_name = get_node_name(node) + is_class = node.type in class_types + + chunk = ASTChunk( + content=content, + content_type=ContentType.FUNCTIONS_CLASSES, + language=lang_name, + path=path, + semantic_names=[node_name] if node_name else [], + parent_context=list(parent_context), + start_line=start_line, + end_line=end_line, + node_type=node.type, + ) + + chunk.docstring = self._metadata_extractor.extract_docstring(content, lang_name) + chunk.signature = self._metadata_extractor.extract_signature(content, lang_name) + + # Extract inheritance via regex + inheritance = self._metadata_extractor.extract_inheritance(content, lang_name) + chunk.extends = inheritance.get('extends', []) + chunk.implements = inheritance.get('implements', []) + chunk.imports = inheritance.get('imports', []) + + # Extract rich AST details directly from this node + self._extract_rich_details_from_node(chunk, node, source_bytes, lang_name) + + chunks.append(chunk) + processed_ranges.add(node_range) + + if is_class and node_name: + for child in node.children: + traverse(child, parent_context + [node_name]) + else: + for child in node.children: + traverse(child, parent_context) + + traverse(tree.root_node, []) + + # Create simplified code + if chunks: + simplified = self._create_simplified_code(text, chunks, lang_name) + if simplified and len(simplified.strip()) > 50: + chunks.append(ASTChunk( + content=simplified, + content_type=ContentType.SIMPLIFIED_CODE, + language=lang_name, + path=path, + start_line=1, + end_line=text.count('\n') + 1, + node_type='simplified', + )) + + return chunks + + def _extract_rich_ast_details( + self, + chunk: ASTChunk, + tree: Any, + captured_node: Any, + lang_name: str + ) -> None: + """ + Extract rich AST details from tree-sitter node by traversing its children. + + This extracts: + - Methods (for classes) + - Properties/fields (for classes) + - Parameters (for functions/methods) + - Return type + - Decorators/annotations + - Called functions/methods + - Referenced types + - Variables + - Type parameters (generics) + """ + source_bytes = chunk.content.encode('utf-8') + + # Find the actual tree-sitter node for this capture + node = self._find_node_at_position( + tree.root_node, + captured_node.start_byte, + captured_node.end_byte + ) + if not node: + return + + # Language-specific node type mappings + node_types = self._get_rich_node_types(lang_name) + + def get_text(n) -> str: + """Get text for a node relative to chunk content.""" + start = n.start_byte - captured_node.start_byte + end = n.end_byte - captured_node.start_byte + if 0 <= start < len(source_bytes) and start < end <= len(source_bytes): + return source_bytes[start:end].decode('utf-8', errors='replace') + return '' + + def extract_identifier(n) -> Optional[str]: + """Extract identifier name from a node.""" + for child in n.children: + if child.type in node_types['identifier']: + return get_text(child) + return None + + def traverse_for_details(n, depth: int = 0): + """Recursively traverse to extract details.""" + if depth > 10: # Prevent infinite recursion + return + + node_type = n.type + + # Extract methods (for classes) + if node_type in node_types['method']: + method_name = extract_identifier(n) + if method_name and method_name not in chunk.methods: + chunk.methods.append(method_name) + + # Extract properties/fields + if node_type in node_types['property']: + prop_name = extract_identifier(n) + if prop_name and prop_name not in chunk.properties: + chunk.properties.append(prop_name) + + # Extract parameters + if node_type in node_types['parameter']: + param_name = extract_identifier(n) + if param_name and param_name not in chunk.parameters: + chunk.parameters.append(param_name) + + # Extract decorators/annotations + if node_type in node_types['decorator']: + dec_text = get_text(n).strip() + if dec_text and dec_text not in chunk.decorators: + # Clean up decorator text + if dec_text.startswith('@'): + dec_text = dec_text[1:] + if '(' in dec_text: + dec_text = dec_text.split('(')[0] + chunk.decorators.append(dec_text) + + # Extract function calls + if node_type in node_types['call']: + call_name = extract_identifier(n) + if call_name and call_name not in chunk.calls: + chunk.calls.append(call_name) + + # Extract type references + if node_type in node_types['type_ref']: + type_text = get_text(n).strip() + if type_text and type_text not in chunk.referenced_types: + # Clean generic params + if '<' in type_text: + type_text = type_text.split('<')[0] + chunk.referenced_types.append(type_text) + + # Extract return type + if node_type in node_types['return_type'] and not chunk.return_type: + chunk.return_type = get_text(n).strip() + + # Extract type parameters (generics) + if node_type in node_types['type_param']: + param_text = get_text(n).strip() + if param_text and param_text not in chunk.type_parameters: + chunk.type_parameters.append(param_text) + + # Extract variables + if node_type in node_types['variable']: + var_name = extract_identifier(n) + if var_name and var_name not in chunk.variables: + chunk.variables.append(var_name) + + # Recurse into children + for child in n.children: + traverse_for_details(child, depth + 1) + + traverse_for_details(node) + + # Limit list sizes to prevent bloat + chunk.methods = chunk.methods[:30] + chunk.properties = chunk.properties[:30] + chunk.parameters = chunk.parameters[:20] + chunk.decorators = chunk.decorators[:10] + chunk.calls = chunk.calls[:50] + chunk.referenced_types = chunk.referenced_types[:30] + chunk.variables = chunk.variables[:30] + chunk.type_parameters = chunk.type_parameters[:10] + + def _find_node_at_position(self, root, start_byte: int, end_byte: int) -> Optional[Any]: + """Find the tree-sitter node at the given byte position.""" + def find(node): + if node.start_byte == start_byte and node.end_byte == end_byte: + return node + for child in node.children: + if child.start_byte <= start_byte and child.end_byte >= end_byte: + result = find(child) + if result: + return result + return None + return find(root) + + def _get_rich_node_types(self, language: str) -> Dict[str, List[str]]: + """Get tree-sitter node types for extracting rich details.""" + # Common patterns across languages + common = { + 'identifier': ['identifier', 'name', 'type_identifier', 'property_identifier'], + 'call': ['call_expression', 'call', 'function_call', 'method_invocation'], + 'type_ref': ['type_identifier', 'generic_type', 'type_annotation', 'type'], + 'type_param': ['type_parameter', 'type_parameters', 'generic_parameter'], + } + + types = { + 'python': { + **common, + 'method': ['function_definition'], + 'property': ['assignment', 'expression_statement'], + 'parameter': ['parameter', 'default_parameter', 'typed_parameter'], + 'decorator': ['decorator'], + 'return_type': ['type'], + 'variable': ['assignment'], + }, + 'java': { + **common, + 'method': ['method_declaration', 'constructor_declaration'], + 'property': ['field_declaration'], + 'parameter': ['formal_parameter', 'spread_parameter'], + 'decorator': ['annotation', 'marker_annotation'], + 'return_type': ['type_identifier', 'generic_type', 'void_type'], + 'variable': ['local_variable_declaration'], + }, + 'javascript': { + **common, + 'method': ['method_definition', 'function_declaration'], + 'property': ['field_definition', 'public_field_definition'], + 'parameter': ['formal_parameters', 'required_parameter'], + 'decorator': ['decorator'], + 'return_type': ['type_annotation'], + 'variable': ['variable_declarator'], + }, + 'typescript': { + **common, + 'method': ['method_definition', 'method_signature', 'function_declaration'], + 'property': ['public_field_definition', 'property_signature'], + 'parameter': ['required_parameter', 'optional_parameter'], + 'decorator': ['decorator'], + 'return_type': ['type_annotation'], + 'variable': ['variable_declarator'], + }, + 'go': { + **common, + 'method': ['method_declaration', 'function_declaration'], + 'property': ['field_declaration'], + 'parameter': ['parameter_declaration'], + 'decorator': [], # Go doesn't have decorators + 'return_type': ['type_identifier', 'pointer_type'], + 'variable': ['short_var_declaration', 'var_declaration'], + }, + 'rust': { + **common, + 'method': ['function_item', 'associated_item'], + 'property': ['field_declaration'], + 'parameter': ['parameter'], + 'decorator': ['attribute_item'], + 'return_type': ['type_identifier', 'generic_type'], + 'variable': ['let_declaration'], + }, + 'c_sharp': { + **common, + 'method': ['method_declaration', 'constructor_declaration'], + 'property': ['property_declaration', 'field_declaration'], + 'parameter': ['parameter'], + 'decorator': ['attribute_list', 'attribute'], + 'return_type': ['predefined_type', 'generic_name'], + 'variable': ['variable_declaration'], + }, + 'php': { + **common, + 'method': ['method_declaration', 'function_definition'], + 'property': ['property_declaration'], + 'parameter': ['simple_parameter'], + 'decorator': ['attribute_list'], + 'return_type': ['named_type', 'union_type'], + 'variable': ['property_declaration', 'simple_variable'], + }, + } + + return types.get(language, { + **common, + 'method': [], + 'property': [], + 'parameter': [], + 'decorator': [], + 'return_type': [], + 'variable': [], + }) + + def _extract_rich_details_from_node( + self, + chunk: ASTChunk, + node: Any, + source_bytes: bytes, + lang_name: str + ) -> None: + """ + Extract rich AST details directly from a tree-sitter node. + Used by traversal-based extraction when we already have the node. + """ + node_types = self._get_rich_node_types(lang_name) + + def get_text(n) -> str: + return source_bytes[n.start_byte:n.end_byte].decode('utf-8', errors='replace') + + def extract_identifier(n) -> Optional[str]: + for child in n.children: + if child.type in node_types['identifier']: + return get_text(child) + return None + + def traverse(n, depth: int = 0): + if depth > 10: + return + + node_type = n.type + + if node_type in node_types['method']: + name = extract_identifier(n) + if name and name not in chunk.methods: + chunk.methods.append(name) + + if node_type in node_types['property']: + name = extract_identifier(n) + if name and name not in chunk.properties: + chunk.properties.append(name) + + if node_type in node_types['parameter']: + name = extract_identifier(n) + if name and name not in chunk.parameters: + chunk.parameters.append(name) + + if node_type in node_types['decorator']: + dec_text = get_text(n).strip() + if dec_text and dec_text not in chunk.decorators: + if dec_text.startswith('@'): + dec_text = dec_text[1:] + if '(' in dec_text: + dec_text = dec_text.split('(')[0] + chunk.decorators.append(dec_text) + + if node_type in node_types['call']: + name = extract_identifier(n) + if name and name not in chunk.calls: + chunk.calls.append(name) + + if node_type in node_types['type_ref']: + type_text = get_text(n).strip() + if type_text and type_text not in chunk.referenced_types: + if '<' in type_text: + type_text = type_text.split('<')[0] + chunk.referenced_types.append(type_text) + + if node_type in node_types['return_type'] and not chunk.return_type: + chunk.return_type = get_text(n).strip() + + if node_type in node_types['type_param']: + param_text = get_text(n).strip() + if param_text and param_text not in chunk.type_parameters: + chunk.type_parameters.append(param_text) + + if node_type in node_types['variable']: + name = extract_identifier(n) + if name and name not in chunk.variables: + chunk.variables.append(name) + + for child in n.children: + traverse(child, depth + 1) + + traverse(node) + + # Limit sizes + chunk.methods = chunk.methods[:30] + chunk.properties = chunk.properties[:30] + chunk.parameters = chunk.parameters[:20] + chunk.decorators = chunk.decorators[:10] + chunk.calls = chunk.calls[:50] + chunk.referenced_types = chunk.referenced_types[:30] + chunk.variables = chunk.variables[:30] + chunk.type_parameters = chunk.type_parameters[:10] + + def _process_chunks( + self, + chunks: List[ASTChunk], + doc: LlamaDocument, + language: Language, + path: str + ) -> List[TextNode]: + """Process AST chunks into TextNodes, handling oversized chunks.""" + nodes = [] + chunk_counter = 0 + + for ast_chunk in chunks: + if len(ast_chunk.content) > self.max_chunk_size: + sub_nodes = self._split_oversized_chunk(ast_chunk, language, doc.metadata, path) + nodes.extend(sub_nodes) + chunk_counter += len(sub_nodes) + else: + metadata = self._build_metadata(ast_chunk, doc.metadata, chunk_counter, len(chunks)) + chunk_id = generate_deterministic_id(path, ast_chunk.content, chunk_counter) + + # Create embedding-enriched text with semantic context + enriched_text = self._create_embedding_text(ast_chunk.content, metadata) + + node = TextNode( + id_=chunk_id, + text=enriched_text, + metadata=metadata + ) + nodes.append(node) + chunk_counter += 1 + + return nodes + + def _split_oversized_chunk( + self, + chunk: ASTChunk, + language: Optional[Language], + base_metadata: Dict[str, Any], + path: str + ) -> List[TextNode]: + """ + Split an oversized chunk using RecursiveCharacterTextSplitter. + + IMPORTANT: Splitting an AST chunk loses semantic integrity. + We try to preserve what we can: + - Parent context and primary name are kept (they're still relevant) + - Detailed lists (methods, properties, calls) are NOT copied to sub-chunks + because they describe the whole unit, not the fragment + - A summary of the original unit is prepended to help embeddings + """ + splitter = self._get_text_splitter(language) if language else self._default_splitter + sub_chunks = splitter.split_text(chunk.content) + + nodes = [] + parent_id = generate_deterministic_id(path, chunk.content, 0) + total_sub = len([s for s in sub_chunks if s and s.strip()]) + + # Build a brief summary of the original semantic unit + # This helps embeddings understand context even in fragments + unit_summary_parts = [] + if chunk.semantic_names: + unit_summary_parts.append(f"{chunk.node_type or 'code'}: {chunk.semantic_names[0]}") + if chunk.extends: + unit_summary_parts.append(f"extends {', '.join(chunk.extends[:3])}") + if chunk.implements: + unit_summary_parts.append(f"implements {', '.join(chunk.implements[:3])}") + if chunk.methods: + unit_summary_parts.append(f"has {len(chunk.methods)} methods") + + unit_summary = " | ".join(unit_summary_parts) if unit_summary_parts else None + + sub_idx = 0 + for i, sub_chunk in enumerate(sub_chunks): + if not sub_chunk or not sub_chunk.strip(): + continue + if len(sub_chunk.strip()) < self.min_chunk_size and total_sub > 1: + continue + + # Build metadata for this fragment + # DO NOT copy detailed lists - they don't apply to fragments + metadata = dict(base_metadata) + metadata['content_type'] = ContentType.OVERSIZED_SPLIT.value + metadata['original_content_type'] = chunk.content_type.value + metadata['parent_chunk_id'] = parent_id + metadata['sub_chunk_index'] = sub_idx + metadata['total_sub_chunks'] = total_sub + metadata['start_line'] = chunk.start_line + metadata['end_line'] = chunk.end_line + + # Keep parent context - still relevant + if chunk.parent_context: + metadata['parent_context'] = chunk.parent_context + metadata['parent_class'] = chunk.parent_context[-1] + + # Keep primary name - this fragment belongs to this unit + if chunk.semantic_names: + metadata['semantic_names'] = chunk.semantic_names[:1] # Just the main name + metadata['primary_name'] = chunk.semantic_names[0] + + # Add note that this is a fragment + metadata['is_fragment'] = True + metadata['fragment_of'] = chunk.semantic_names[0] if chunk.semantic_names else None + + # For embedding: prepend fragment context + if unit_summary: + fragment_header = f"[Fragment {sub_idx + 1}/{total_sub} of {unit_summary}]" + enriched_text = f"{fragment_header}\n\n{sub_chunk}" + else: + enriched_text = sub_chunk + + chunk_id = generate_deterministic_id(path, sub_chunk, sub_idx) + nodes.append(TextNode(id_=chunk_id, text=enriched_text, metadata=metadata)) + sub_idx += 1 + + # Log when splitting happens - it's a signal the chunk_size might need adjustment + if nodes: + logger.info( + f"Split oversized {chunk.node_type or 'chunk'} " + f"'{chunk.semantic_names[0] if chunk.semantic_names else 'unknown'}' " + f"({len(chunk.content)} chars) into {len(nodes)} fragments" + ) + + return nodes + + def _split_fallback( + self, + doc: LlamaDocument, + language: Optional[Language] = None + ) -> List[TextNode]: + """Fallback splitting using RecursiveCharacterTextSplitter.""" + text = doc.text + path = doc.metadata.get('path', 'unknown') + + if not text or not text.strip(): + return [] + + splitter = self._get_text_splitter(language) if language else self._default_splitter + chunks = splitter.split_text(text) + + nodes = [] + lang_str = doc.metadata.get('language', 'text') + text_offset = 0 + + for i, chunk in enumerate(chunks): + if not chunk or not chunk.strip(): + continue + if len(chunk.strip()) < self.min_chunk_size and len(chunks) > 1: + continue + if len(chunk) > 30000: + chunk = chunk[:30000] + + # Calculate line numbers + start_line = text[:text_offset].count('\n') + 1 if text_offset > 0 else 1 + chunk_pos = text.find(chunk, text_offset) + if chunk_pos >= 0: + text_offset = chunk_pos + len(chunk) + end_line = start_line + chunk.count('\n') + + metadata = dict(doc.metadata) + metadata['content_type'] = ContentType.FALLBACK.value + metadata['chunk_index'] = i + metadata['total_chunks'] = len(chunks) + metadata['start_line'] = start_line + metadata['end_line'] = end_line + + # Extract names via regex + names = self._metadata_extractor.extract_names_from_content(chunk, lang_str) + if names: + metadata['semantic_names'] = names + metadata['primary_name'] = names[0] + + # Extract inheritance + inheritance = self._metadata_extractor.extract_inheritance(chunk, lang_str) + if inheritance.get('extends'): + metadata['extends'] = inheritance['extends'] + metadata['parent_types'] = inheritance['extends'] + if inheritance.get('implements'): + metadata['implements'] = inheritance['implements'] + if inheritance.get('imports'): + metadata['imports'] = inheritance['imports'] + + # Create embedding-enriched text with semantic context + enriched_text = self._create_embedding_text(chunk, metadata) + + chunk_id = generate_deterministic_id(path, chunk, i) + nodes.append(TextNode(id_=chunk_id, text=enriched_text, metadata=metadata)) + + return nodes + + def _build_metadata( + self, + chunk: ASTChunk, + base_metadata: Dict[str, Any], + chunk_index: int, + total_chunks: int + ) -> Dict[str, Any]: + """Build metadata dictionary from ASTChunk.""" + metadata = dict(base_metadata) + + metadata['content_type'] = chunk.content_type.value + metadata['node_type'] = chunk.node_type + metadata['chunk_index'] = chunk_index + metadata['total_chunks'] = total_chunks + metadata['start_line'] = chunk.start_line + metadata['end_line'] = chunk.end_line + + if chunk.parent_context: + metadata['parent_context'] = chunk.parent_context + metadata['parent_class'] = chunk.parent_context[-1] + metadata['full_path'] = '.'.join(chunk.parent_context + chunk.semantic_names[:1]) + + if chunk.semantic_names: + metadata['semantic_names'] = chunk.semantic_names + metadata['primary_name'] = chunk.semantic_names[0] + + if chunk.docstring: + metadata['docstring'] = chunk.docstring[:500] + + if chunk.signature: + metadata['signature'] = chunk.signature + + if chunk.extends: + metadata['extends'] = chunk.extends + metadata['parent_types'] = chunk.extends + + if chunk.implements: + metadata['implements'] = chunk.implements + + if chunk.imports: + metadata['imports'] = chunk.imports + + if chunk.namespace: + metadata['namespace'] = chunk.namespace + + # --- RICH AST METADATA --- + + if chunk.methods: + metadata['methods'] = chunk.methods + + if chunk.properties: + metadata['properties'] = chunk.properties + + if chunk.parameters: + metadata['parameters'] = chunk.parameters + + if chunk.return_type: + metadata['return_type'] = chunk.return_type + + if chunk.decorators: + metadata['decorators'] = chunk.decorators + + if chunk.modifiers: + metadata['modifiers'] = chunk.modifiers + + if chunk.calls: + metadata['calls'] = chunk.calls + + if chunk.referenced_types: + metadata['referenced_types'] = chunk.referenced_types + + if chunk.variables: + metadata['variables'] = chunk.variables + + if chunk.constants: + metadata['constants'] = chunk.constants + + if chunk.type_parameters: + metadata['type_parameters'] = chunk.type_parameters + + return metadata + + def _create_embedding_text(self, content: str, metadata: Dict[str, Any]) -> str: + """ + Create embedding-optimized text by prepending concise semantic context. + + Design principles: + 1. Keep it SHORT - long headers can skew embeddings for small code chunks + 2. Avoid redundancy - don't repeat info that's obvious from the code + 3. Clean paths - strip commit hashes and archive prefixes + 4. Add VALUE - include info that helps semantic matching + + What we include (selectively): + - Clean file path (without commit/archive prefixes) + - Parent context (for nested structures - very valuable) + - Extends/implements (inheritance is critical for understanding) + - Docstring (helps semantic matching) + - For CLASSES: method count (helps identify scope) + - For METHODS: skip redundant method list + """ + if not self.enrich_embedding_text: + return content + + context_parts = [] + + # Clean file path - remove commit hash prefixes and archive structure + path = metadata.get('path', '') + if path: + path = self._clean_path(path) + context_parts.append(f"File: {path}") + + # Parent context - valuable for nested structures + parent_context = metadata.get('parent_context', []) + if parent_context: + context_parts.append(f"In: {'.'.join(parent_context)}") + + # Clean namespace - strip keyword if present + namespace = metadata.get('namespace', '') + if namespace: + ns_clean = namespace.replace('namespace ', '').replace('package ', '').strip().rstrip(';') + if ns_clean: + context_parts.append(f"Namespace: {ns_clean}") + + # Type relationships - very valuable for understanding code structure + extends = metadata.get('extends', []) + implements = metadata.get('implements', []) + if extends: + context_parts.append(f"Extends: {', '.join(extends[:3])}") + if implements: + context_parts.append(f"Implements: {', '.join(implements[:3])}") + + # For CLASSES: show method/property counts (helps understand scope) + # For METHODS/FUNCTIONS: skip - it's redundant + node_type = metadata.get('node_type', '') + is_container = node_type in ('class', 'interface', 'struct', 'trait', 'enum', 'impl') + + if is_container: + methods = metadata.get('methods', []) + properties = metadata.get('properties', []) + if methods and len(methods) > 1: + # Only show if there are multiple methods + context_parts.append(f"Methods({len(methods)}): {', '.join(methods[:8])}") + if properties and len(properties) > 1: + context_parts.append(f"Fields({len(properties)}): {', '.join(properties[:5])}") + + # Docstring - valuable for semantic matching + docstring = metadata.get('docstring', '') + if docstring: + # Take just the first sentence or 100 chars + brief = docstring.split('.')[0][:100].strip() + if brief: + context_parts.append(f"Desc: {brief}") + + # Build final text - only if we have meaningful context + if context_parts: + context_header = " | ".join(context_parts) + return f"[{context_header}]\n\n{content}" + + return content + + def _clean_path(self, path: str) -> str: + """ + Clean file path for embedding text. + + Removes: + - Commit hash prefixes (e.g., 'owner-repo-abc123def/') + - Archive extraction paths + - Redundant path components + """ + if not path: + return path + + # Split by '/' and look for src/, lib/, app/ etc as anchor points + parts = path.split('/') + + # Common source directory markers + source_markers = {'src', 'lib', 'app', 'source', 'main', 'test', 'tests', 'pkg', 'cmd', 'internal'} + + # Find the first source marker and start from there + for i, part in enumerate(parts): + if part.lower() in source_markers: + return '/'.join(parts[i:]) + + # If no marker found but path has commit-hash-like prefix (40 hex chars or similar) + if parts and len(parts) > 1: + first_part = parts[0] + # Check if first part looks like "owner-repo-commithash" pattern + if '-' in first_part and len(first_part) > 40: + # Skip the first part + return '/'.join(parts[1:]) + + return path + + def _get_text_splitter(self, language: Language) -> RecursiveCharacterTextSplitter: + """Get language-specific text splitter.""" + if language not in self._splitter_cache: + try: + self._splitter_cache[language] = RecursiveCharacterTextSplitter.from_language( + language=language, + chunk_size=self.max_chunk_size, + chunk_overlap=self.chunk_overlap, + ) + except Exception: + self._splitter_cache[language] = self._default_splitter + return self._splitter_cache[language] + + def _create_simplified_code( + self, + source_code: str, + chunks: List[ASTChunk], + language: str + ) -> str: + """Create simplified code with placeholders for extracted chunks.""" + semantic_chunks = [c for c in chunks if c.content_type == ContentType.FUNCTIONS_CLASSES] + if not semantic_chunks: + return source_code + + sorted_chunks = sorted( + semantic_chunks, + key=lambda x: source_code.find(x.content), + reverse=True + ) + + result = source_code + comment_prefix = self._metadata_extractor.get_comment_prefix(language) + + for chunk in sorted_chunks: + pos = result.find(chunk.content) + if pos == -1: + continue + + first_line = chunk.content.split('\n')[0].strip() + if len(first_line) > 60: + first_line = first_line[:60] + '...' + + breadcrumb = "" + if chunk.parent_context: + breadcrumb = f" (in {'.'.join(chunk.parent_context)})" + + placeholder = f"{comment_prefix} Code for: {first_line}{breadcrumb}\n" + result = result[:pos] + placeholder + result[pos + len(chunk.content):] + + return result.strip() + + def _parse_type_list(self, text: str) -> List[str]: + """Parse a comma-separated list of types.""" + if not text: + return [] + + text = text.strip().strip('()[]') + + # Remove keywords + for kw in ('extends', 'implements', 'with', ':'): + text = text.replace(kw, ' ') + + types = [] + for part in text.split(','): + name = part.strip() + if '<' in name: + name = name.split('<')[0].strip() + if '(' in name: + name = name.split('(')[0].strip() + if name: + types.append(name) + + return types + + def _get_semantic_node_types(self, language: str) -> Dict[str, List[str]]: + """Get semantic node types for manual traversal fallback.""" + types = { + 'python': { + 'class': ['class_definition'], + 'function': ['function_definition'], + }, + 'java': { + 'class': ['class_declaration', 'interface_declaration', 'enum_declaration'], + 'function': ['method_declaration', 'constructor_declaration'], + }, + 'javascript': { + 'class': ['class_declaration'], + 'function': ['function_declaration', 'method_definition', 'arrow_function'], + }, + 'typescript': { + 'class': ['class_declaration', 'interface_declaration'], + 'function': ['function_declaration', 'method_definition', 'arrow_function'], + }, + 'go': { + 'class': ['type_declaration'], + 'function': ['function_declaration', 'method_declaration'], + }, + 'rust': { + 'class': ['struct_item', 'impl_item', 'trait_item', 'enum_item'], + 'function': ['function_item'], + }, + 'c_sharp': { + 'class': ['class_declaration', 'interface_declaration', 'struct_declaration'], + 'function': ['method_declaration', 'constructor_declaration'], + }, + 'php': { + 'class': ['class_declaration', 'interface_declaration', 'trait_declaration'], + 'function': ['function_definition', 'method_declaration'], + }, + } + return types.get(language, {'class': [], 'function': []}) + + @staticmethod + def get_supported_languages() -> List[str]: + """Return list of languages with AST support.""" + return list(LANGUAGE_TO_TREESITTER.values()) + + @staticmethod + def is_ast_supported(path: str) -> bool: + """Check if AST parsing is supported for a file.""" + return is_ast_supported(path) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/tree_parser.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/tree_parser.py new file mode 100644 index 00000000..da378a38 --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/core/splitter/tree_parser.py @@ -0,0 +1,129 @@ +""" +Tree-sitter parser wrapper with caching and language loading. + +Handles dynamic loading of tree-sitter language modules using the new API (v0.23+). +""" + +import logging +from typing import Dict, Any, Optional + +from .languages import TREESITTER_MODULES + +logger = logging.getLogger(__name__) + + +class TreeSitterParser: + """ + Wrapper for tree-sitter parser with language caching. + + Uses the new tree-sitter API (v0.23+) with individual language packages. + """ + + def __init__(self): + self._language_cache: Dict[str, Any] = {} + self._available: Optional[bool] = None + + def is_available(self) -> bool: + """Check if tree-sitter is available and working.""" + if self._available is None: + try: + from tree_sitter import Parser, Language + import tree_sitter_python as tspython + + py_language = Language(tspython.language()) + parser = Parser(py_language) + parser.parse(b"def test(): pass") + + self._available = True + logger.info("tree-sitter is available and working") + except ImportError as e: + logger.warning(f"tree-sitter not installed: {e}") + self._available = False + except Exception as e: + logger.warning(f"tree-sitter error: {type(e).__name__}: {e}") + self._available = False + return self._available + + def get_language(self, lang_name: str) -> Optional[Any]: + """ + Get tree-sitter Language object for a language name. + + Args: + lang_name: Tree-sitter language name (e.g., 'python', 'java', 'php') + + Returns: + tree_sitter.Language object or None if unavailable + """ + if lang_name in self._language_cache: + return self._language_cache[lang_name] + + if not self.is_available(): + return None + + try: + from tree_sitter import Language + + lang_info = TREESITTER_MODULES.get(lang_name) + if not lang_info: + logger.debug(f"No tree-sitter module mapping for '{lang_name}'") + return None + + module_name, func_name = lang_info + + import importlib + lang_module = importlib.import_module(module_name) + + lang_func = getattr(lang_module, func_name, None) + if not lang_func: + logger.debug(f"Module {module_name} has no {func_name} function") + return None + + language = Language(lang_func()) + self._language_cache[lang_name] = language + return language + + except Exception as e: + logger.debug(f"Could not load tree-sitter language '{lang_name}': {e}") + return None + + def parse(self, source_code: str, lang_name: str) -> Optional[Any]: + """ + Parse source code and return the AST tree. + + Args: + source_code: Source code string + lang_name: Tree-sitter language name + + Returns: + tree_sitter.Tree object or None if parsing failed + """ + language = self.get_language(lang_name) + if not language: + return None + + try: + from tree_sitter import Parser + + parser = Parser(language) + tree = parser.parse(bytes(source_code, "utf8")) + return tree + + except Exception as e: + logger.warning(f"Failed to parse code with tree-sitter ({lang_name}): {e}") + return None + + def clear_cache(self): + """Clear the language cache.""" + self._language_cache.clear() + + +# Global singleton instance +_parser_instance: Optional[TreeSitterParser] = None + + +def get_parser() -> TreeSitterParser: + """Get the global TreeSitterParser instance.""" + global _parser_instance + if _parser_instance is None: + _parser_instance = TreeSitterParser() + return _parser_instance diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/models/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/models/__init__.py index dd0b41b6..02c5632f 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/models/__init__.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/models/__init__.py @@ -1,6 +1,6 @@ """Models and configuration for RAG pipeline""" -from .config import RAGConfig, DocumentMetadata, IndexStats +from .config import RAGConfig, IndexStats -__all__ = ["RAGConfig", "DocumentMetadata", "IndexStats"] +__all__ = ["RAGConfig", "IndexStats"] diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/models/config.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/models/config.py index bde5e2bf..2cffc100 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/models/config.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/models/config.py @@ -1,11 +1,14 @@ import os -from typing import Optional +from typing import Optional, Literal from pydantic import BaseModel, Field, field_validator, model_validator import logging from dotenv import load_dotenv logger = logging.getLogger(__name__) +# Embedding provider types +EmbeddingProvider = Literal["ollama", "openrouter"] + # Known embedding model dimensions EMBEDDING_MODEL_DIMENSIONS = { # OpenAI models @@ -18,9 +21,17 @@ "cohere/embed-multilingual-v3.0": 1024, "voyage/voyage-large-2": 1536, "voyage/voyage-code-2": 1536, - # Alibaba models (often 4096) + # Alibaba/Qwen models "qwen/qwen-embedding": 4096, + "qwen/qwen3-embedding-0.6b": 1024, "qwen/qwen3-embedding-8b": 4096, + # Ollama local models (same Qwen models) + "qwen3-embedding-0.6b": 1024, + "qwen3-embedding:0.6b": 1024, + "snowflake-arctic-embed": 1024, + "nomic-embed-text": 768, + "mxbai-embed-large": 1024, + "all-minilm": 384, # Default fallback "default": 1536, } @@ -40,14 +51,21 @@ def get_embedding_dim_for_model(model: str) -> int: class RAGConfig(BaseModel): """Configuration for RAG pipeline""" - load_dotenv() + load_dotenv(interpolate=False) # Qdrant for vector storage qdrant_url: str = Field(default_factory=lambda: os.getenv("QDRANT_URL", "http://localhost:6333")) qdrant_collection_prefix: str = Field(default_factory=lambda: os.getenv("QDRANT_COLLECTION_PREFIX", "codecrow")) - # OpenRouter for embeddings + # Embedding provider selection: "ollama" (local) or "openrouter" (cloud) + embedding_provider: str = Field(default_factory=lambda: os.getenv("EMBEDDING_PROVIDER", "ollama")) + + # Ollama configuration (local embeddings - default) + ollama_base_url: str = Field(default_factory=lambda: os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")) + ollama_model: str = Field(default_factory=lambda: os.getenv("OLLAMA_EMBEDDING_MODEL", "qwen3-embedding:0.6b")) + + # OpenRouter configuration (cloud embeddings - optional) openrouter_api_key: str = Field(default_factory=lambda: os.getenv("OPENROUTER_API_KEY", "")) - openrouter_model: str = Field(default_factory=lambda: os.getenv("OPENROUTER_MODEL", "openai/text-embedding-3-small")) + openrouter_model: str = Field(default_factory=lambda: os.getenv("OPENROUTER_MODEL", "qwen/qwen3-embedding-8b")) openrouter_base_url: str = Field(default="https://openrouter.ai/api/v1") # Embedding dimensions - auto-detected from model or set via env var @@ -58,26 +76,39 @@ class RAGConfig(BaseModel): def set_embedding_dim_from_model(self) -> 'RAGConfig': """Auto-detect embedding dimension from model if not explicitly set.""" if self.embedding_dim == 0: - self.embedding_dim = get_embedding_dim_for_model(self.openrouter_model) - logger.info(f"Auto-detected embedding dimension {self.embedding_dim} for model {self.openrouter_model}") + # Use the model for the selected provider + model = self.ollama_model if self.embedding_provider == "ollama" else self.openrouter_model + self.embedding_dim = get_embedding_dim_for_model(model) + logger.info(f"Auto-detected embedding dimension {self.embedding_dim} for model {model}") else: logger.info(f"Using configured embedding dimension: {self.embedding_dim}") return self - @field_validator('openrouter_api_key') - @classmethod - def validate_api_key(cls, v: str) -> str: - if not v or v.strip() == "": - logger.error("OPENROUTER_API_KEY is not set or is empty!") - logger.error("Please set the OPENROUTER_API_KEY environment variable") - raise ValueError("OPENROUTER_API_KEY is required but not set") - logger.info(f"OpenRouter API key loaded: {v[:10]}...{v[-4:]}") - return v - - chunk_size: int = Field(default=800) + @model_validator(mode='after') + def validate_provider_config(self) -> 'RAGConfig': + """Validate that the selected provider has required configuration.""" + if self.embedding_provider == "openrouter": + if not self.openrouter_api_key or self.openrouter_api_key.strip() == "": + logger.error("OPENROUTER_API_KEY is not set but embedding_provider is 'openrouter'!") + raise ValueError("OPENROUTER_API_KEY is required when using OpenRouter provider") + logger.info(f"Using OpenRouter embeddings with model: {self.openrouter_model}") + logger.info(f"OpenRouter API key loaded: {self.openrouter_api_key[:10]}...{self.openrouter_api_key[-4:]}") + elif self.embedding_provider == "ollama": + logger.info(f"Using Ollama local embeddings with model: {self.ollama_model}") + logger.info(f"Ollama base URL: {self.ollama_base_url}") + else: + logger.warning(f"Unknown embedding provider '{self.embedding_provider}', defaulting to 'ollama'") + self.embedding_provider = "ollama" + return self + + # Chunk size for code files + # text-embedding-3-small supports ~8191 tokens (~32K chars) + # 8000 chars keeps most semantic units (classes, functions) intact + chunk_size: int = Field(default=8000) chunk_overlap: int = Field(default=200) - text_chunk_size: int = Field(default=1000) + # Text chunk size for non-code files (markdown, docs) + text_chunk_size: int = Field(default=2000) text_chunk_overlap: int = Field(default=200) base_index_namespace: str = Field(default="code_rag") @@ -120,19 +151,6 @@ def validate_api_key(cls, v: str) -> str: max_chunks_per_index: int = Field(default_factory=lambda: int(os.getenv("RAG_MAX_CHUNKS_PER_INDEX", "10000"))) max_files_per_index: int = Field(default_factory=lambda: int(os.getenv("RAG_MAX_FILES_PER_INDEX", "500"))) - -class DocumentMetadata(BaseModel): - """Metadata for indexed documents""" - workspace: str - project: str - branch: str - path: str - commit: str - language: str - filetype: str - chunk_index: Optional[int] = None - - class IndexStats(BaseModel): """Statistics about an index""" namespace: str diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/models/scoring_config.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/models/scoring_config.py new file mode 100644 index 00000000..f5735c5a --- /dev/null +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/models/scoring_config.py @@ -0,0 +1,257 @@ +""" +Scoring configuration for RAG query result reranking. + +Provides configurable boost factors and priority patterns that can be +overridden via environment variables. +""" + +import os +from typing import Dict, List +from pydantic import BaseModel, Field +import logging + +logger = logging.getLogger(__name__) + + +def _parse_list_env(env_var: str, default: List[str]) -> List[str]: + """Parse comma-separated environment variable into list.""" + value = os.getenv(env_var) + if not value: + return default + return [item.strip() for item in value.split(',') if item.strip()] + + +def _parse_float_env(env_var: str, default: float) -> float: + """Parse float from environment variable.""" + value = os.getenv(env_var) + if not value: + return default + try: + return float(value) + except ValueError: + logger.warning(f"Invalid float value for {env_var}: {value}, using default {default}") + return default + + +class ContentTypeBoost(BaseModel): + """Boost factors for different content types from AST parsing.""" + + functions_classes: float = Field( + default_factory=lambda: _parse_float_env("RAG_BOOST_FUNCTIONS_CLASSES", 1.2), + description="Boost for full function/class definitions (highest value)" + ) + fallback: float = Field( + default_factory=lambda: _parse_float_env("RAG_BOOST_FALLBACK", 1.0), + description="Boost for regex-based split chunks" + ) + oversized_split: float = Field( + default_factory=lambda: _parse_float_env("RAG_BOOST_OVERSIZED", 0.95), + description="Boost for large chunks that were split" + ) + simplified_code: float = Field( + default_factory=lambda: _parse_float_env("RAG_BOOST_SIMPLIFIED", 0.7), + description="Boost for code with placeholders (context only)" + ) + + def get(self, content_type: str) -> float: + """Get boost factor for a content type.""" + return getattr(self, content_type, 1.0) + + +class FilePriorityPatterns(BaseModel): + """File path patterns for priority-based boosting.""" + + high: List[str] = Field( + default_factory=lambda: _parse_list_env( + "RAG_HIGH_PRIORITY_PATTERNS", + ['service', 'controller', 'handler', 'api', 'core', 'auth', 'security', + 'permission', 'repository', 'dao', 'migration'] + ), + description="Patterns for high-priority files (1.3x boost)" + ) + + medium: List[str] = Field( + default_factory=lambda: _parse_list_env( + "RAG_MEDIUM_PRIORITY_PATTERNS", + ['model', 'entity', 'dto', 'schema', 'util', 'helper', 'common', + 'shared', 'component', 'hook', 'client', 'integration'] + ), + description="Patterns for medium-priority files (1.1x boost)" + ) + + low: List[str] = Field( + default_factory=lambda: _parse_list_env( + "RAG_LOW_PRIORITY_PATTERNS", + ['test', 'spec', 'config', 'mock', 'fixture', 'stub'] + ), + description="Patterns for low-priority files (0.8x penalty)" + ) + + high_boost: float = Field( + default_factory=lambda: _parse_float_env("RAG_HIGH_PRIORITY_BOOST", 1.3) + ) + medium_boost: float = Field( + default_factory=lambda: _parse_float_env("RAG_MEDIUM_PRIORITY_BOOST", 1.1) + ) + low_boost: float = Field( + default_factory=lambda: _parse_float_env("RAG_LOW_PRIORITY_BOOST", 0.8) + ) + + def get_priority(self, file_path: str) -> tuple: + """ + Get priority level and boost factor for a file path. + + Uses word-boundary matching to avoid false positives like 'test' matching 'latest.py'. + Patterns are matched against path segments (directories and filename). + + Returns: + Tuple of (priority_name, boost_factor) + """ + import re + + path_lower = file_path.lower() + # Extract path segments for word-boundary matching + segments = re.split(r'[/\\]', path_lower) + # Also consider filename without extension + filename = segments[-1] if segments else '' + name_without_ext = filename.rsplit('.', 1)[0] if '.' in filename else filename + + def pattern_matches(patterns: List[str]) -> bool: + for p in patterns: + # Check if pattern matches as a complete segment + if p in segments: + return True + # Check if pattern matches at word boundaries in filename + # E.g., 'test' matches 'test_utils.py' or 'UserServiceTest.java' but not 'latest.py' + if re.search(rf'\b{re.escape(p)}\b', name_without_ext): + return True + # Also check directory names for patterns like 'tests/', '__tests__/' + for seg in segments[:-1]: + if re.search(rf'\b{re.escape(p)}\b', seg): + return True + return False + + if pattern_matches(self.high): + return ('HIGH', self.high_boost) + elif pattern_matches(self.medium): + return ('MEDIUM', self.medium_boost) + elif pattern_matches(self.low): + return ('LOW', self.low_boost) + else: + return ('MEDIUM', 1.0) + + +class MetadataBonus(BaseModel): + """Bonus multipliers for metadata presence.""" + + semantic_names: float = Field( + default_factory=lambda: _parse_float_env("RAG_BONUS_SEMANTIC_NAMES", 1.1), + description="Bonus for chunks with extracted semantic names" + ) + docstring: float = Field( + default_factory=lambda: _parse_float_env("RAG_BONUS_DOCSTRING", 1.05), + description="Bonus for chunks with docstrings" + ) + signature: float = Field( + default_factory=lambda: _parse_float_env("RAG_BONUS_SIGNATURE", 1.02), + description="Bonus for chunks with function signatures" + ) + + +class ScoringConfig(BaseModel): + """ + Complete scoring configuration for RAG query reranking. + + All values can be overridden via environment variables: + - RAG_BOOST_FUNCTIONS_CLASSES, RAG_BOOST_FALLBACK, etc. + - RAG_HIGH_PRIORITY_PATTERNS (comma-separated) + - RAG_HIGH_PRIORITY_BOOST, RAG_MEDIUM_PRIORITY_BOOST, etc. + - RAG_BONUS_SEMANTIC_NAMES, RAG_BONUS_DOCSTRING, etc. + + Usage: + config = ScoringConfig() + boost = config.content_type_boost.get('functions_classes') + priority, boost = config.file_priority.get_priority('/src/UserService.java') + """ + + content_type_boost: ContentTypeBoost = Field(default_factory=ContentTypeBoost) + file_priority: FilePriorityPatterns = Field(default_factory=FilePriorityPatterns) + metadata_bonus: MetadataBonus = Field(default_factory=MetadataBonus) + + # Score thresholds + min_relevance_score: float = Field( + default_factory=lambda: _parse_float_env("RAG_MIN_RELEVANCE_SCORE", 0.7), + description="Minimum score threshold for results" + ) + + max_score_cap: float = Field( + default_factory=lambda: _parse_float_env("RAG_MAX_SCORE_CAP", 1.0), + description="Maximum score cap after boosting" + ) + + def calculate_boosted_score( + self, + base_score: float, + file_path: str, + content_type: str, + has_semantic_names: bool = False, + has_docstring: bool = False, + has_signature: bool = False + ) -> tuple: + """ + Calculate final boosted score for a result. + + Args: + base_score: Original similarity score + file_path: File path of the chunk + content_type: Content type (functions_classes, fallback, etc.) + has_semantic_names: Whether chunk has semantic names + has_docstring: Whether chunk has docstring + has_signature: Whether chunk has signature + + Returns: + Tuple of (boosted_score, priority_level) + """ + score = base_score + + # File priority boost + priority, priority_boost = self.file_priority.get_priority(file_path) + score *= priority_boost + + # Content type boost + content_boost = self.content_type_boost.get(content_type) + score *= content_boost + + # Metadata bonuses + if has_semantic_names: + score *= self.metadata_bonus.semantic_names + if has_docstring: + score *= self.metadata_bonus.docstring + if has_signature: + score *= self.metadata_bonus.signature + + # Cap the score + score = min(score, self.max_score_cap) + + return (score, priority) + + +# Global singleton +_scoring_config: ScoringConfig | None = None + + +def get_scoring_config() -> ScoringConfig: + """Get the global ScoringConfig instance.""" + global _scoring_config + if _scoring_config is None: + _scoring_config = ScoringConfig() + logger.info("ScoringConfig initialized with:") + logger.info(f" High priority patterns: {_scoring_config.file_priority.high[:5]}...") + logger.info(f" Content type boosts: functions_classes={_scoring_config.content_type_boost.functions_classes}") + return _scoring_config + + +def reset_scoring_config(): + """Reset the global config (useful for testing).""" + global _scoring_config + _scoring_config = None diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/services/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/services/__init__.py index 1bb93040..88077d7a 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/services/__init__.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/services/__init__.py @@ -1,7 +1,6 @@ """Services for querying and webhook integration""" from .query_service import RAGQueryService -from .webhook_integration import WebhookIntegration -__all__ = ["RAGQueryService", "WebhookIntegration"] +__all__ = ["RAGQueryService"] diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/services/query_service.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/services/query_service.py index a8dcf3f5..06768098 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/services/query_service.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/services/query_service.py @@ -1,43 +1,19 @@ from typing import List, Dict, Optional import logging -import os from llama_index.core import VectorStoreIndex from llama_index.vector_stores.qdrant import QdrantVectorStore from qdrant_client import QdrantClient -from qdrant_client.http.models import Filter, FieldCondition, MatchValue, MatchAny +from qdrant_client.http.models import Filter, FieldCondition, MatchValue, MatchAny, MatchText from ..models.config import RAGConfig -from ..utils.utils import make_namespace, make_project_namespace -from ..core.openrouter_embedding import OpenRouterEmbedding +from ..models.scoring_config import get_scoring_config +from ..utils.utils import make_project_namespace +from ..core.embedding_factory import create_embedding_model, get_embedding_model_info from ..models.instructions import InstructionType, format_query logger = logging.getLogger(__name__) -# File priority patterns for smart RAG -HIGH_PRIORITY_PATTERNS = [ - 'service', 'controller', 'handler', 'api', 'core', 'auth', 'security', - 'permission', 'repository', 'dao', 'migration' -] - -MEDIUM_PRIORITY_PATTERNS = [ - 'model', 'entity', 'dto', 'schema', 'util', 'helper', 'common', - 'shared', 'component', 'hook', 'client', 'integration' -] - -LOW_PRIORITY_PATTERNS = [ - 'test', 'spec', 'config', 'mock', 'fixture', 'stub' -] - -# Content type priorities for AST-based chunks -# functions_classes are more valuable than simplified_code (placeholders) -CONTENT_TYPE_BOOST = { - 'functions_classes': 1.2, # Full function/class definitions - highest value - 'fallback': 1.0, # Regex-based split - normal value - 'oversized_split': 0.95, # Large chunks that were split - slightly lower - 'simplified_code': 0.7, # Code with placeholders - lower value (context only) -} - class RAGQueryService: """Service for querying RAG indices using Qdrant. @@ -52,14 +28,11 @@ def __init__(self, config: RAGConfig): # Qdrant client self.qdrant_client = QdrantClient(url=config.qdrant_url) - # Embedding model - self.embed_model = OpenRouterEmbedding( - api_key=config.openrouter_api_key, - model=config.openrouter_model, - api_base=config.openrouter_base_url, - timeout=60.0, - max_retries=3 - ) + # Embedding model (supports Ollama and OpenRouter via factory) + embed_info = get_embedding_model_info(config) + logger.info(f"QueryService using embedding provider: {embed_info['provider']} ({embed_info['type']})") + + self.embed_model = create_embedding_model(config) def _collection_or_alias_exists(self, name: str) -> bool: """Check if a collection or alias with the given name exists.""" @@ -82,11 +55,6 @@ def _get_project_collection_name(self, workspace: str, project: str) -> str: namespace = make_project_namespace(workspace, project) return f"{self.config.qdrant_collection_prefix}_{namespace}" - def _get_collection_name(self, workspace: str, project: str, branch: str) -> str: - """Generate collection name (legacy - kept for backward compatibility)""" - namespace = make_namespace(workspace, project, branch) - return f"{self.config.qdrant_collection_prefix}_{namespace}" - def _dedupe_by_branch_priority( self, results: List[Dict], @@ -95,45 +63,64 @@ def _dedupe_by_branch_priority( ) -> List[Dict]: """Deduplicate results by file path, preferring target branch version. - When same file exists in multiple branches, keep only one version: - - Prefer target_branch version (it's the latest) - - Fall back to base_branch version if target doesn't have it + When same file exists in multiple branches, keep only the TARGET branch version. + This ensures we review the NEW code, not the OLD code. + + Strategy: + 1. First pass: collect all paths that exist in target branch + 2. Second pass: for each result, include it only if: + - It's from target branch, OR + - Its path doesn't exist in target branch (cross-file reference from base) - This preserves cross-file relationships while avoiding duplicates. + This ensures: + - Changed files are always from target branch (the PR's new code) + - Related files from base branch are included only if they don't exist in target """ if not results: return results - # Group by path + chunk position (approximate by content hash) - grouped = {} + # Step 1: Find all paths that exist in target branch + target_branch_paths = set() + for result in results: + metadata = result.get('metadata', {}) + branch = metadata.get('branch', '') + if branch == target_branch: + path = metadata.get('path', metadata.get('file_path', '')) + target_branch_paths.add(path) + + logger.debug(f"Target branch '{target_branch}' has {len(target_branch_paths)} unique paths") + + # Step 2: Filter results - target branch wins for same path + deduped = [] + seen_chunks = set() # Track (path, chunk_identity) to avoid exact duplicates for result in results: metadata = result.get('metadata', {}) path = metadata.get('path', metadata.get('file_path', '')) branch = metadata.get('branch', '') - # Create a key based on path and approximate content position - # Using text hash to distinguish different chunks from same file - text_hash = hash(result.get('text', '')[:200]) # First 200 chars for identity - key = f"{path}:{text_hash}" + # Create chunk identity (path + start of content) + chunk_id = f"{path}:{branch}:{hash(result.get('text', '')[:100])}" - if key not in grouped: - grouped[key] = result - else: - existing_branch = grouped[key].get('metadata', {}).get('branch', '') - - # Prefer target branch, then base branch, then whatever has higher score - if branch == target_branch and existing_branch != target_branch: - grouped[key] = result - elif (branch == base_branch and - existing_branch != target_branch and - existing_branch != base_branch): - grouped[key] = result - elif result['score'] > grouped[key]['score'] and branch == existing_branch: - # Same branch, keep higher score - grouped[key] = result - - return list(grouped.values()) + if chunk_id in seen_chunks: + continue + seen_chunks.add(chunk_id) + + # Include if: + # 1. It's from target branch (always include), OR + # 2. Path doesn't exist in target branch (cross-file reference from base) + if branch == target_branch: + deduped.append(result) + elif path not in target_branch_paths: + # This file only exists in base branch - include for cross-file context + deduped.append(result) + # else: skip - file exists in target branch, use that version instead + + skipped_count = len(results) - len(deduped) + if skipped_count > 0: + logger.info(f"Branch priority: kept {len(deduped)} results, skipped {skipped_count} base branch duplicates") + + return deduped def semantic_search( self, @@ -288,6 +275,403 @@ def _get_fallback_branch(self, workspace: str, project: str, requested_branch: s return None + def get_deterministic_context( + self, + workspace: str, + project: str, + branches: List[str], + file_paths: List[str], + limit_per_file: int = 10 + ) -> Dict: + """ + Get context using DETERMINISTIC metadata-based retrieval. + + Leverages ALL tree-sitter metadata extracted during indexing: + - semantic_names: function/method/class names + - primary_name: main identifier + - parent_class: containing class + - full_path: qualified name (e.g., "Data.getConfigData") + - imports: import statements + - extends: parent classes/interfaces + - namespace: package/namespace + - node_type: method_declaration, class_definition, etc. + + Multi-step process: + 1. Query chunks for changed file_paths + 2. Extract metadata (identifiers, parent classes, namespaces, imports) + 3. Find related definitions by: + a) primary_name match (definitions of used identifiers) + b) parent_class match (other methods in same class) + c) namespace match (related code in same package) + + NO LANGUAGE-SPECIFIC PARSING NEEDED - tree-sitter already did that! + Same input always produces same output (deterministic). + + Args: + workspace: VCS workspace + project: Project name + branches: Branches to search (target + base for PRs) + file_paths: Changed file paths from diff + limit_per_file: Max chunks per file + + Returns: + Dict with chunks grouped by retrieval type and rich metadata + """ + collection_name = self._get_project_collection_name(workspace, project) + + if not self._collection_or_alias_exists(collection_name): + logger.warning(f"Collection {collection_name} does not exist") + return {"chunks": [], "changed_files": {}, "related_definitions": {}, + "class_context": {}, "namespace_context": {}, + "_metadata": {"error": "collection_not_found"}} + + logger.info(f"Deterministic context: files={file_paths[:5]}, branches={branches}") + + def _apply_branch_priority(points: list, target: str, existing_target_paths: set) -> list: + """Filter points to prioritize target branch. + + For each unique path: + - If path exists in target branch, keep only target branch version + - If path only in base branch, keep it (cross-file reference) + """ + if not target or len(branches) == 1: + return points + + # Group by path + by_path = {} + for p in points: + path = p.payload.get("path", "") + if path not in by_path: + by_path[path] = [] + by_path[path].append(p) + + # Select best version per path + result = [] + for path, path_points in by_path.items(): + has_target = any(p.payload.get("branch") == target for p in path_points) + if has_target: + # Keep only target branch for this path + result.extend([p for p in path_points if p.payload.get("branch") == target]) + elif path not in existing_target_paths: + # Path doesn't exist in target - keep base branch version + result.extend(path_points) + # else: skip - path exists in target but these results are from base + + return result + + all_chunks = [] + changed_files_chunks = {} + related_definitions = {} + class_context = {} # Other methods in same classes + namespace_context = {} # Related code in same namespaces + + # Metadata to collect from changed files + identifiers_to_find = set() + parent_classes = set() + namespaces = set() + imports_raw = set() + extends_raw = set() + + # Track changed file paths for deduplication + changed_file_paths = set() + seen_texts = set() + + # Build branch filter - NOTE: branches[0] is the target branch (has priority) + target_branch = branches[0] if branches else None + branch_filter = ( + FieldCondition(key="branch", match=MatchValue(value=branches[0])) + if len(branches) == 1 + else FieldCondition(key="branch", match=MatchAny(any=branches)) + ) + + # Track which paths exist in target branch (for priority filtering) + target_branch_paths = set() + + # ========== STEP 1: Get chunks from changed files ========== + for file_path in file_paths: + try: + normalized_path = file_path.lstrip("/") + filename = normalized_path.rsplit("/", 1)[-1] if "/" in normalized_path else normalized_path + + # Try exact path match + results, _ = self.qdrant_client.scroll( + collection_name=collection_name, + scroll_filter=Filter( + must=[ + branch_filter, + FieldCondition(key="path", match=MatchValue(value=normalized_path)) + ] + ), + limit=limit_per_file * len(branches), # Get more to account for multiple branches + with_payload=True, + with_vectors=False + ) + + # Fallback: try with filename only (handles path prefix mismatches) + if not results: + results, _ = self.qdrant_client.scroll( + collection_name=collection_name, + scroll_filter=Filter( + must=[ + branch_filter, + FieldCondition(key="path", match=MatchText(text=filename)) + ] + ), + limit=limit_per_file * len(branches), + with_payload=True, + with_vectors=False + ) + + # Apply branch priority: if file exists in target branch, only keep target branch version + if target_branch and len(branches) > 1: + # Check if any result is from target branch + has_target = any(p.payload.get("branch") == target_branch for p in results) + if has_target: + # Keep only target branch results for this file + results = [p for p in results if p.payload.get("branch") == target_branch] + logger.debug(f"Branch priority: keeping target branch '{target_branch}' for {normalized_path}") + + # Apply limit after filtering + results = results[:limit_per_file] + + chunks_for_file = [] + for point in results: + payload = point.payload + text = payload.get("text", payload.get("_node_content", "")) + + if text in seen_texts: + continue + seen_texts.add(text) + + # Track which paths exist in target branch + if payload.get("branch") == target_branch: + target_branch_paths.add(payload.get("path", "")) + + chunk = { + "text": text, + "metadata": {k: v for k, v in payload.items() if k not in ("text", "_node_content")}, + "score": 1.0, + "_match_type": "changed_file", + "_matched_on": file_path + } + chunks_for_file.append(chunk) + all_chunks.append(chunk) + changed_file_paths.add(payload.get("path", "")) + + # Extract ALL tree-sitter metadata for step 2-4 + if isinstance(payload.get("semantic_names"), list): + identifiers_to_find.update(payload["semantic_names"]) + + if payload.get("primary_name"): + identifiers_to_find.add(payload["primary_name"]) + + if payload.get("parent_class"): + parent_classes.add(payload["parent_class"]) + + if payload.get("namespace"): + namespaces.add(payload["namespace"]) + + if isinstance(payload.get("imports"), list): + for imp in payload["imports"]: + # Extract class name from import statement + # "use Magento\Store\Model\ScopeInterface;" -> "ScopeInterface" + if isinstance(imp, str): + parts = imp.replace(";", "").split("\\") + if parts: + imports_raw.add(parts[-1].strip()) + + if isinstance(payload.get("extends"), list): + extends_raw.update(payload["extends"]) + + if payload.get("parent_class"): + extends_raw.add(payload["parent_class"]) + + changed_files_chunks[file_path] = chunks_for_file + + except Exception as e: + logger.warning(f"Error querying file '{file_path}': {e}") + + logger.info(f"Step 1: {len(all_chunks)} chunks from changed files. " + f"Extracted: {len(identifiers_to_find)} identifiers, " + f"{len(parent_classes)} parent_classes, {len(namespaces)} namespaces, " + f"{len(imports_raw)} imports, {len(extends_raw)} extends") + + # ========== STEP 2: Find definitions by primary_name ========== + # Find where identifiers/imports/extends are DEFINED + all_to_find = identifiers_to_find | imports_raw | extends_raw + if all_to_find: + try: + batch = list(all_to_find)[:100] + results, _ = self.qdrant_client.scroll( + collection_name=collection_name, + scroll_filter=Filter( + must=[ + branch_filter, + FieldCondition(key="primary_name", match=MatchAny(any=batch)) + ] + ), + limit=200 * len(branches), # Get more to account for multiple branches + with_payload=True, + with_vectors=False + ) + + # Apply branch priority filtering + results = _apply_branch_priority(results, target_branch, target_branch_paths) + + for point in results: + payload = point.payload + if payload.get("path") in changed_file_paths: + continue + + text = payload.get("text", payload.get("_node_content", "")) + if text in seen_texts: + continue + seen_texts.add(text) + + primary_name = payload.get("primary_name", "") + chunk = { + "text": text, + "metadata": {k: v for k, v in payload.items() if k not in ("text", "_node_content")}, + "score": 0.95, + "_match_type": "definition", + "_matched_on": primary_name + } + all_chunks.append(chunk) + + if primary_name not in related_definitions: + related_definitions[primary_name] = [] + related_definitions[primary_name].append(chunk) + + logger.info(f"Step 2: Found {len(related_definitions)} definitions by primary_name") + + except Exception as e: + logger.warning(f"Error in primary_name query: {e}") + + # ========== STEP 3: Find other methods in same parent_class ========== + # If we're changing a method in class "Data", find other methods of "Data" + if parent_classes: + try: + batch = list(parent_classes)[:20] + results, _ = self.qdrant_client.scroll( + collection_name=collection_name, + scroll_filter=Filter( + must=[ + branch_filter, + FieldCondition(key="parent_class", match=MatchAny(any=batch)) + ] + ), + limit=100 * len(branches), # Get more to account for multiple branches + with_payload=True, + with_vectors=False + ) + + # Apply branch priority filtering + results = _apply_branch_priority(results, target_branch, target_branch_paths) + + for point in results: + payload = point.payload + if payload.get("path") in changed_file_paths: + continue + + text = payload.get("text", payload.get("_node_content", "")) + if text in seen_texts: + continue + seen_texts.add(text) + + parent_class = payload.get("parent_class", "") + chunk = { + "text": text, + "metadata": {k: v for k, v in payload.items() if k not in ("text", "_node_content")}, + "score": 0.85, + "_match_type": "class_context", + "_matched_on": parent_class + } + all_chunks.append(chunk) + + if parent_class not in class_context: + class_context[parent_class] = [] + class_context[parent_class].append(chunk) + + logger.info(f"Step 3: Found {sum(len(v) for v in class_context.values())} class context chunks") + + except Exception as e: + logger.warning(f"Error in parent_class query: {e}") + + # ========== STEP 4: Find related code in same namespace ========== + # Lower priority - only get a few for broader context + if namespaces: + try: + batch = list(namespaces)[:10] + results, _ = self.qdrant_client.scroll( + collection_name=collection_name, + scroll_filter=Filter( + must=[ + branch_filter, + FieldCondition(key="namespace", match=MatchAny(any=batch)) + ] + ), + limit=30 * len(branches), # Get more to account for multiple branches + with_payload=True, + with_vectors=False + ) + + # Apply branch priority filtering + results = _apply_branch_priority(results, target_branch, target_branch_paths) + + for point in results: + payload = point.payload + if payload.get("path") in changed_file_paths: + continue + + text = payload.get("text", payload.get("_node_content", "")) + if text in seen_texts: + continue + seen_texts.add(text) + + namespace = payload.get("namespace", "") + chunk = { + "text": text, + "metadata": {k: v for k, v in payload.items() if k not in ("text", "_node_content")}, + "score": 0.75, + "_match_type": "namespace_context", + "_matched_on": namespace + } + all_chunks.append(chunk) + + if namespace not in namespace_context: + namespace_context[namespace] = [] + namespace_context[namespace].append(chunk) + + logger.info(f"Step 4: Found {sum(len(v) for v in namespace_context.values())} namespace context chunks") + + except Exception as e: + logger.warning(f"Error in namespace query: {e}") + + logger.info(f"Deterministic context complete: {len(all_chunks)} total chunks " + f"(changed: {sum(len(v) for v in changed_files_chunks.values())}, " + f"definitions: {sum(len(v) for v in related_definitions.values())}, " + f"class_ctx: {sum(len(v) for v in class_context.values())}, " + f"ns_ctx: {sum(len(v) for v in namespace_context.values())})") + + return { + "chunks": all_chunks, + "changed_files": changed_files_chunks, + "related_definitions": related_definitions, + "class_context": class_context, + "namespace_context": namespace_context, + "_metadata": { + "branches_searched": branches, + "target_branch": target_branch, + "files_requested": file_paths, + "identifiers_extracted": list(identifiers_to_find)[:30], + "parent_classes_found": list(parent_classes), + "namespaces_found": list(namespaces), + "imports_extracted": list(imports_raw)[:30], + "extends_extracted": list(extends_raw)[:20], + "target_branch_paths_found": len(target_branch_paths) + } + } + def get_context_for_pr( self, workspace: str, @@ -301,7 +685,8 @@ def get_context_for_pr( enable_priority_reranking: bool = True, min_relevance_score: float = 0.7, base_branch: Optional[str] = None, - deleted_files: Optional[List[str]] = None + deleted_files: Optional[List[str]] = None, + exclude_pr_files: Optional[List[str]] = None ) -> Dict: """ Get relevant context for PR review using Smart RAG with multi-branch support. @@ -313,9 +698,14 @@ def get_context_for_pr( branch: Target branch (the PR's source branch) base_branch: Base branch (the PR's target, e.g., 'main'). If None, uses fallback logic. deleted_files: Files that were deleted in target branch (excluded from results) + exclude_pr_files: Files indexed separately as PR data (excluded to avoid duplication) """ diff_snippets = diff_snippets or [] deleted_files = deleted_files or [] + exclude_pr_files = exclude_pr_files or [] + + # Combine exclusion lists: deleted files + PR-indexed files + all_excluded_paths = list(set(deleted_files + exclude_pr_files)) # Determine branches to search branches_to_search = [branch] @@ -356,11 +746,15 @@ def get_context_for_pr( diff_snippets=diff_snippets, changed_files=changed_files ) + + logger.info(f"Generated {len(queries)} queries for PR context") + for i, (q_text, q_weight, q_top_k, q_type) in enumerate(queries): + logger.info(f" Query {i+1}: weight={q_weight}, top_k={q_top_k}, text='{q_text[:80]}...'") all_results = [] # 2. Execute queries with multi-branch search - for q_text, q_weight, q_top_k, q_instruction_type in queries: + for i, (q_text, q_weight, q_top_k, q_instruction_type) in enumerate(queries): if not q_text.strip(): continue @@ -371,8 +765,10 @@ def get_context_for_pr( branches=branches_to_search, top_k=q_top_k, instruction_type=q_instruction_type, - excluded_paths=deleted_files + excluded_paths=all_excluded_paths ) + + logger.info(f"Query {i+1}/{len(queries)} returned {len(results)} results") for r in results: r["_query_weight"] = q_weight @@ -419,9 +815,12 @@ def get_context_for_pr( if "path" in result["metadata"]: related_files.add(result["metadata"]["path"]) - logger.info( - f"Smart RAG: Final context has {len(relevant_code)} chunks " - f"from {len(related_files)} files across {len(branches_to_search)} branches") + # Log top results for debugging + logger.info(f"Smart RAG: Final context has {len(relevant_code)} chunks from {len(related_files)} files") + for i, r in enumerate(relevant_code[:5]): + path = r["metadata"].get("path", "unknown") + primary_name = r["metadata"].get("primary_name", "N/A") + logger.info(f" Chunk {i+1}: score={r['score']:.3f}, name={primary_name}, path=...{path[-60:]}") result = { "relevant_code": relevant_code, @@ -444,6 +843,7 @@ def _decompose_queries( """ from collections import defaultdict import os + import re queries = [] @@ -486,14 +886,33 @@ def _decompose_queries( queries.append((q, 0.8, 5, InstructionType.LOGIC)) # C. Snippet Queries (Low Level) - Weight 1.2 (High precision) - for snippet in diff_snippets[:3]: - # Clean snippet: remove +/- markers, take first few lines - lines = [l.strip() for l in snippet.split('\n') if l.strip() and not l.startswith(('+', '-'))] + # Use actual changed code for semantic matching (not just context lines) + for snippet in diff_snippets[:5]: + # Extract meaningful code from diff - INCLUDE changed lines (+/-), they ARE the code + lines = [] + for line in snippet.split('\n'): + stripped = line.strip() + if not stripped: + continue + # Skip diff headers but keep actual code (including +/- prefixed lines) + if stripped.startswith(('diff --git', '---', '+++', '@@', 'index ')): + continue + # Remove the +/- prefix but keep the code content + if stripped.startswith('+') or stripped.startswith('-'): + code_line = stripped[1:].strip() + if code_line and len(code_line) > 3: # Skip empty/trivial lines + lines.append(code_line) + elif stripped: + lines.append(stripped) + if lines: - # Join first 2-3 significant lines - clean_snippet = " ".join(lines[:3]) - if len(clean_snippet) > 10: - queries.append((clean_snippet, 1.2, 5, InstructionType.DEPENDENCY)) + # Join significant lines (function names, method calls, etc.) + clean_snippet = " ".join(lines[:5]) + if len(clean_snippet) > 15: + queries.append((clean_snippet, 1.2, 8, InstructionType.DEPENDENCY)) + + # Log the generated queries for debugging + logger.debug(f"Decomposed into {len(queries)} queries: {[(q[0][:50], q[1]) for q in queries]}") return queries @@ -501,11 +920,12 @@ def _merge_and_rank_results(self, results: List[Dict], min_score_threshold: floa """ Deduplicate matches and filter by relevance score with priority-based reranking. - Applies three types of boosting: + Uses ScoringConfig for configurable boosting factors: 1. File path priority (service/controller vs test/config) 2. Content type priority (functions_classes vs simplified_code) 3. Semantic name bonus (chunks with extracted function/class names) """ + scoring_config = get_scoring_config() grouped = {} # Deduplicate by file_path + content hash @@ -521,43 +941,28 @@ def _merge_and_rank_results(self, results: List[Dict], min_score_threshold: floa unique_results = list(grouped.values()) - # Apply multi-factor score boosting + # Apply multi-factor score boosting using ScoringConfig for result in unique_results: metadata = result.get('metadata', {}) - file_path = metadata.get('path', metadata.get('file_path', '')).lower() + file_path = metadata.get('path', metadata.get('file_path', '')) content_type = metadata.get('content_type', 'fallback') semantic_names = metadata.get('semantic_names', []) + has_docstring = bool(metadata.get('docstring')) + has_signature = bool(metadata.get('signature')) + + boosted_score, priority = scoring_config.calculate_boosted_score( + base_score=result['score'], + file_path=file_path, + content_type=content_type, + has_semantic_names=bool(semantic_names), + has_docstring=has_docstring, + has_signature=has_signature + ) - base_score = result['score'] - - # 1. File path priority boosting - if any(p in file_path for p in HIGH_PRIORITY_PATTERNS): - base_score *= 1.3 - result['_priority'] = 'HIGH' - elif any(p in file_path for p in MEDIUM_PRIORITY_PATTERNS): - base_score *= 1.1 - result['_priority'] = 'MEDIUM' - elif any(p in file_path for p in LOW_PRIORITY_PATTERNS): - base_score *= 0.8 # Penalize test/config files - result['_priority'] = 'LOW' - else: - result['_priority'] = 'MEDIUM' - - # 2. Content type boosting (AST-based metadata) - content_boost = CONTENT_TYPE_BOOST.get(content_type, 1.0) - base_score *= content_boost + result['score'] = boosted_score + result['_priority'] = priority result['_content_type'] = content_type - - # 3. Semantic name bonus - chunks with extracted names are more valuable - if semantic_names: - base_score *= 1.1 # 10% bonus for having semantic names - result['_has_semantic_names'] = True - - # 4. Docstring bonus - chunks with docstrings provide better context - if metadata.get('docstring'): - base_score *= 1.05 # 5% bonus for having docstring - - result['score'] = min(1.0, base_score) + result['_has_semantic_names'] = bool(semantic_names) # Filter by threshold filtered = [r for r in unique_results if r['score'] >= min_score_threshold] diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/services/webhook_integration.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/services/webhook_integration.py deleted file mode 100644 index 4183d4b0..00000000 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/services/webhook_integration.py +++ /dev/null @@ -1,253 +0,0 @@ -""" -Integration service to connect RAG pipeline with CodeCrow pipeline agent -""" - -import logging -from typing import List, Dict, Optional -from pathlib import Path -import tempfile -import shutil - -from ..models.config import RAGConfig -from ..core.index_manager import RAGIndexManager -from .query_service import RAGQueryService - -logger = logging.getLogger(__name__) - - -class WebhookIntegration: - """Handles webhook events and manages RAG indexing""" - - def __init__(self, config: RAGConfig): - self.config = config - self.index_manager = RAGIndexManager(config) - self.query_service = RAGQueryService(config) - - def handle_pr_created( - self, - workspace: str, - project: str, - branch: str, - commit: str, - repo_clone_path: str, - pr_description: Optional[str] = None - ) -> Dict: - """ - Handle PR created event - index entire repository - - Args: - workspace: Workspace slug - project: Project key - branch: Branch name - commit: Commit SHA - repo_clone_path: Path to cloned repository - pr_description: PR description (optional) - - Returns: - Dict with indexing stats and initial context - """ - logger.info(f"Handling PR created for {workspace}/{project}/{branch}") - - try: - stats = self.index_manager.index_repository( - repo_path=repo_clone_path, - workspace=workspace, - project=project, - branch=branch, - commit=commit - ) - - logger.info(f"Indexed {stats.document_count} documents, {stats.chunk_count} chunks") - - return { - "success": True, - "stats": stats.dict(), - "message": f"Indexed {stats.document_count} documents" - } - - except Exception as e: - logger.error(f"Error handling PR created: {e}") - return { - "success": False, - "error": str(e) - } - - def handle_pr_updated( - self, - workspace: str, - project: str, - branch: str, - commit: str, - changed_files: List[str], - repo_clone_path: str - ) -> Dict: - """ - Handle PR updated event - incrementally update changed files - - Args: - workspace: Workspace slug - project: Project key - branch: Branch name - commit: Commit SHA - changed_files: List of changed file paths - repo_clone_path: Path to cloned repository - - Returns: - Dict with update stats - """ - logger.info(f"Handling PR updated for {workspace}/{project}/{branch}") - logger.info(f"Changed files: {len(changed_files)}") - - try: - added_modified = [] - deleted = [] - - for file_path in changed_files: - full_path = Path(repo_clone_path) / file_path - if full_path.exists(): - added_modified.append(file_path) - else: - deleted.append(file_path) - - if deleted: - self.index_manager.delete_files( - file_paths=deleted, - workspace=workspace, - project=project, - branch=branch - ) - logger.info(f"Deleted {len(deleted)} files from index") - - if added_modified: - stats = self.index_manager.update_files( - file_paths=added_modified, - repo_base=repo_clone_path, - workspace=workspace, - project=project, - branch=branch, - commit=commit - ) - logger.info(f"Updated {len(added_modified)} files") - else: - stats = self.index_manager._get_index_stats(workspace, project, branch) - - return { - "success": True, - "stats": stats.dict(), - "files_updated": len(added_modified), - "files_deleted": len(deleted) - } - - except Exception as e: - logger.error(f"Error handling PR updated: {e}") - return { - "success": False, - "error": str(e) - } - - def get_pr_analysis_context( - self, - workspace: str, - project: str, - branch: str, - changed_files: List[str], - pr_description: Optional[str] = None, - top_k: int = 10 - ) -> str: - """ - Get relevant context for PR analysis - - Args: - workspace: Workspace slug - project: Project key - branch: Branch name - changed_files: List of changed file paths - pr_description: PR description (optional) - top_k: Number of top results to return - - Returns: - Formatted context string for LLM - """ - logger.info(f"Getting PR analysis context for {workspace}/{project}/{branch}") - - try: - context = self.query_service.get_context_for_pr( - workspace=workspace, - project=project, - branch=branch, - changed_files=changed_files, - pr_description=pr_description, - top_k=top_k - ) - - return context - - except Exception as e: - logger.error(f"Error getting PR context: {e}") - return f"Error retrieving context: {str(e)}" - - def handle_branch_deleted( - self, - workspace: str, - project: str, - branch: str - ) -> Dict: - """ - Handle branch deleted event - clean up index - - Args: - workspace: Workspace slug - project: Project key - branch: Branch name - - Returns: - Dict with deletion result - """ - logger.info(f"Handling branch deleted for {workspace}/{project}/{branch}") - - try: - self.index_manager.delete_index(workspace, project, branch) - - return { - "success": True, - "message": f"Deleted index for {workspace}/{project}/{branch}" - } - - except Exception as e: - logger.error(f"Error handling branch deleted: {e}") - return { - "success": False, - "error": str(e) - } - - def get_file_context( - self, - workspace: str, - project: str, - branch: str, - file_path: str, - top_k: int = 5 - ) -> str: - """ - Get context related to a specific file - - Args: - workspace: Workspace slug - project: Project key - branch: Branch name - file_path: File path - top_k: Number of results - - Returns: - Formatted context string - """ - context = self.query_service.get_context_for_files( - workspace=workspace, - project=project, - branch=branch, - file_paths=[file_path], - top_k=top_k - ) - - return context - diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/__init__.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/__init__.py index d3432ec3..8610e100 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/__init__.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/__init__.py @@ -5,7 +5,6 @@ make_namespace, should_exclude_file, is_binary_file, - is_code_file, LANGUAGE_MAP ) diff --git a/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/utils.py b/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/utils.py index 4b965cb2..2373c91a 100644 --- a/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/utils.py +++ b/python-ecosystem/rag-pipeline/src/rag_pipeline/utils/utils.py @@ -69,6 +69,53 @@ def make_project_namespace(workspace: str, project: str) -> str: return f"{workspace}__{project}".replace("/", "_").replace(".", "_").lower() +def clean_archive_path(path: str) -> str: + """ + Clean archive root prefix from file paths. + + Bitbucket and other VCS archives often create a root folder like: + - 'owner-repo-commitHash/' (Bitbucket) + - 'repo-branch/' (GitHub) + + This function strips that prefix to get clean paths like 'src/file.php'. + + Args: + path: File path potentially with archive prefix + + Returns: + Clean path without archive prefix + """ + if not path: + return path + + parts = Path(path).parts + if len(parts) <= 1: + return path + + first_part = parts[0] + + # Common source directory markers - if first part is one of these, path is already clean + source_markers = {'src', 'lib', 'app', 'source', 'main', 'test', 'tests', + 'pkg', 'cmd', 'internal', 'bin', 'scripts', 'docs'} + if first_part.lower() in source_markers: + return path + + # Check if first part looks like archive root: + # - Contains hyphens (owner-repo-commit pattern) + # - Or is very long (40+ chars for commit hash) + # - Or matches pattern like 'name-hexstring' + looks_like_archive = ( + '-' in first_part and len(first_part) > 20 or # owner-repo-commit + len(first_part) >= 40 or # Just commit hash + (first_part.count('-') >= 2 and any(c.isdigit() for c in first_part)) # Has digits and multiple hyphens + ) + + if looks_like_archive: + return '/'.join(parts[1:]) + + return path + + def should_exclude_file(path: str, excluded_patterns: list[str]) -> bool: """Check if file should be excluded based on patterns. @@ -144,11 +191,4 @@ def is_binary_file(file_path: Path) -> bool: return True return False except Exception: - return True - - -def is_code_file(language: str) -> bool: - """Check if the language represents code (vs text/config)""" - non_code = {'text', 'markdown', 'rst', 'json', 'xml', 'yaml', 'toml', 'ini', 'config'} - return language not in non_code - + return True \ No newline at end of file diff --git a/python-ecosystem/rag-pipeline/tests/test_rag.py b/python-ecosystem/rag-pipeline/tests/test_rag.py index 4be64146..28a6a775 100644 --- a/python-ecosystem/rag-pipeline/tests/test_rag.py +++ b/python-ecosystem/rag-pipeline/tests/test_rag.py @@ -4,7 +4,7 @@ import pytest from pathlib import Path -from rag_pipeline.models.config import RAGConfig, DocumentMetadata +from rag_pipeline.models.config import RAGConfig from rag_pipeline.utils.utils import ( detect_language_from_path, make_namespace, @@ -60,30 +60,12 @@ def test_config_defaults(): """Test default configuration""" config = RAGConfig() - assert config.chunk_size == 800 + assert config.chunk_size == 8000 # Increased to fit most semantic units assert config.chunk_overlap == 200 - assert config.text_chunk_size == 1000 + assert config.text_chunk_size == 2000 assert config.retrieval_top_k == 10 assert config.similarity_threshold == 0.7 - -def test_document_metadata(): - """Test document metadata model""" - metadata = DocumentMetadata( - workspace="codecrow", - project="demo", - branch="main", - path="src/main.py", - commit="abc123", - language="python", - filetype="py" - ) - - assert metadata.workspace == "codecrow" - assert metadata.language == "python" - assert metadata.chunk_index is None - - if __name__ == "__main__": pytest.main([__file__, "-v"])