Merge pull request #73 from openrewrite/final-touches-ai-code-search

add datatable to debug accuracy and only check method invocations in …

src/main/java/io/moderne/ai/research/FindCodeThatResembles.java

@@ -18,6 +18,7 @@
 import io.moderne.ai.AgentGenerativeModelClient;
 import io.moderne.ai.EmbeddingModelClient;
 import io.moderne.ai.RelatedModelClient;
+import io.moderne.ai.table.CodeSearch;
 import io.moderne.ai.table.EmbeddingPerformance;
 import lombok.EqualsAndHashCode;
 import lombok.RequiredArgsConstructor;
@@ -29,6 +30,7 @@
 import org.openrewrite.java.MethodMatcher;
 import org.openrewrite.java.search.UsesMethod;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaSourceFile;
 import org.openrewrite.marker.SearchResult;
 
 import java.time.Duration;
@@ -54,6 +56,7 @@ public class FindCodeThatResembles extends ScanningRecipe<FindCodeThatResembles.
             example = "1000")
     int k;
 
+    transient CodeSearch codeSearchTable = new CodeSearch(this);
 
     transient EmbeddingPerformance performance = new EmbeddingPerformance(this);
 
@@ -185,6 +188,11 @@ public J.CompilationUnit visitCompilationUnit(J.CompilationUnit cu, ExecutionCon
 
             @Override
             public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
+
+                if (!getLanguage().equals("java")) {
+                    return super.visitMethodInvocation(method, ctx);
+                }
+
                 boolean matches = false;
                 for (MethodMatcher methodMatcher : methodMatchers) {
                     if (methodMatcher.matches(method)) {
@@ -208,11 +216,25 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                 }
                 int resultEmbeddingModels = related.isRelated();
                 boolean result;
+                boolean calledGenerativeModel = false;
                 if (resultEmbeddingModels == 0) {
                     result = AgentGenerativeModelClient.getInstance().isRelated(resembles, method.printTrimmed(getCursor()), 0.5932);
+                    calledGenerativeModel = true;
                 } else {
                     result = resultEmbeddingModels == 1;
                 }
+
+                // Populate data table for debugging model's accuracy
+                JavaSourceFile javaSourceFile = getCursor().firstEnclosing(JavaSourceFile.class);
+                String source = javaSourceFile.getSourcePath().toString();
+                codeSearchTable.insertRow(ctx, new CodeSearch.Row(
+                        source,
+                        method.printTrimmed(getCursor()),
+                        resembles,
+                        resultEmbeddingModels,
+                        calledGenerativeModel ? ( result ? 1 : -1) : 0
+                ));
+
                 return result ?
                         SearchResult.found(method) :
                         super.visitMethodInvocation(method, ctx);

src/main/java/io/moderne/ai/table/CodeSearch.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2021 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.moderne.ai.table;
+
+import lombok.Value;
+import org.openrewrite.Column;
+import org.openrewrite.DataTable;
+import org.openrewrite.Recipe;
+
+
+public class CodeSearch extends DataTable<CodeSearch.Row> {
+
+    public CodeSearch(Recipe recipe) {
+        super(recipe,
+                "Code Search",
+                "Searches for method invocations that resemble a natural language query.");
+    }
+
+    @Value
+    public static class Row {
+        @Column(displayName = "Source",
+                description = "Source")
+        String source;
+
+        @Column(displayName = "Method",
+                description = "Method invocation")
+        String method;
+
+        @Column(displayName = "Query",
+                description = "Natural language query")
+        String query;
+
+        @Column(displayName = "Result of first models",
+                description = "First two embeddings models result," +
+                        " where -1 means negative match, 0 means unsure, and 1 means positive match.")
+        int resultEmbedding;
+
+        @Column(displayName = "Result of second model",
+                description = "Second generative model's result," +
+                        " where -1 means negative match and 1 means positive match. " +
+                        "If the model was never queried, then the result is 0.")
+        int resultGenerative;
+    }
+}