Enable direct RAG

Author: johnsonr

db/admin.cypher

@@ -13,9 +13,16 @@ RETURN count(*) AS droppedIndexes
 
 
 
-CALL db.index.fulltext.queryNodes(`embabel-content-fulltext-index`, 'xy')
+CALL db.index.vector.queryNodes(`embabel-content-index`, 10, $embeddings)
 YIELD node AS chunk, score
   WHERE score >= $similarityThreshold
+RETURN chunk.text AS text, chunk.id AS id,
+       score
+  ORDER BY score DESC
+
+CALL db.index.fulltext.queryNodes('embabel-content-fulltext-index', 'the PromptRunner interface')
+YIELD node AS chunk, score
+  WHERE score >= .8
 RETURN chunk.text AS text, chunk.id AS id,
        score
   ORDER BY score DESC

db/clean.cypher

@@ -0,0 +1,13 @@
+// Delete data
+// You have been warned!
+
+MATCH (n)
+DETACH DELETE n;
+
+DROP INDEX `embabel-content-fulltext-index`;
+
+DROP INDEX `embabel-content-index`;
+
+DROP INDEX `embabel-entity-fulltext-index`;
+
+DROP INDEX `embabel-entity-index`;

src/main/java/com/embabel/guide/GuideAgentBot.java

@@ -9,6 +9,7 @@
 import com.embabel.agent.api.common.SomeOf;
 import com.embabel.agent.core.AgentPlatform;
 import com.embabel.agent.rag.pipeline.event.RagPipelineEvent;
+import com.embabel.agent.rag.tools.DualShotConfig;
 import com.embabel.chat.AssistantMessage;
 import com.embabel.chat.Chatbot;
 import com.embabel.chat.Conversation;
@@ -19,6 +20,7 @@
 import org.springframework.lang.Nullable;
 
 import javax.validation.constraints.Null;
+import java.time.Duration;
 import java.util.Collections;
 
 record ConversationOver(String why) {
@@ -55,11 +57,14 @@ ChatbotReturn respond(
                 .ai()
                 .withLlm(guideData.guideConfig().llm())
                 .withReferences(guideData.references())
-                .withRag(guideData.ragOptions().withListener(e -> {
-                    if (e instanceof RagPipelineEvent rpe) {
-                        context.updateProgress(rpe.getDescription());
-                    }
-                }))
+                .withRag(guideData.ragOptions()
+                        .withDesiredMaxLatency(Duration.ofMinutes(10))
+                        .withDualShot(new DualShotConfig(100))
+                        .withListener(e -> {
+                            if (e instanceof RagPipelineEvent rpe) {
+                                context.updateProgress(rpe.getDescription());
+                            }
+                        }))
                 .withTemplate("guide_system")
                 .respondWithSystemPrompt(conversation,
                         guideData.templateModel(Collections.singletonMap("user",

src/main/java/com/embabel/guide/GuideData.java

@@ -108,4 +108,5 @@ public Map<String, Object> templateModel(Map<String, Object> extras) {
         n.put("persona", guideConfig.persona());
         return n;
     }
+
 }

src/main/java/com/embabel/guide/GuideShell.java

@@ -1,18 +1,24 @@
 package com.embabel.guide;
 
+import com.embabel.agent.api.common.OperationContext;
 import com.embabel.agent.event.logging.personality.severance.LumonColorPalette;
+import com.embabel.agent.rag.RagRequest;
+import com.embabel.agent.rag.SimpleRagResponseFormatter;
 import com.embabel.agent.shell.TerminalServices;
 import com.embabel.chat.Chatbot;
 import org.springframework.shell.standard.ShellComponent;
 import org.springframework.shell.standard.ShellMethod;
+import org.springframework.shell.standard.ShellOption;
 
 import java.nio.file.Path;
+import java.time.Duration;
 
 @ShellComponent
 public record GuideShell(
         TerminalServices terminalServices,
         Chatbot chatbot,
-        GuideData guideData) {
+        GuideData guideData,
+        OperationContext context) {
 
     @ShellMethod("talk to docs")
     public String talk() {
@@ -43,4 +49,20 @@ public String provision() {
     public String chunks() {
         return guideData.count() + " chunks in the database";
     }
+
+    @ShellMethod("rag query")
+    public String rag(
+            @ShellOption String query,
+            @ShellOption(defaultValue = "10") int topK,
+            @ShellOption(defaultValue = "0.2") double similarityThreshold) {
+        var raw = context.ai().rag().search(
+                RagRequest.query(query)
+                        .withTopK(topK)
+                        .withSimilarityThreshold(similarityThreshold)
+                        .withDesiredMaxLatency(Duration.ofMinutes(5))
+        );
+        var ragResponseFormatter = SimpleRagResponseFormatter.INSTANCE;
+        return "RAG response: \n" + ragResponseFormatter.format(raw);
+    }
+
 }

src/main/resources/application.properties

@@ -9,26 +9,37 @@ embabel.agent.rag.neo.uri=${NEO4J_URI:bolt://localhost:7687}
 embabel.agent.rag.neo.username=${NEO4J_USERNAME:neo4j}
 embabel.agent.rag.neo.password=${NEO4J_PASSWORD:brahmsian}
 
+embabel.agent.rag.neo.max-chunk-size=1500
+embabel.agent.rag.neo.overlap-size=200
+
 # Packages that OGM should scan for entities
 #embabel.agent.rag.neo.packages=
 
-#embabel.discord.token=${EMBABEL_TEAM_HELPER_DISCORD_TOKEN:null}
+embabel.discord.token=${EMBABEL_TEAM_HELPER_DISCORD_TOKEN:null}
 
 
-embabel.agent.rag.compression-llm.model=${guide.default-llm}
-embabel.agent.rag.reranking-llm.model=${guide.default-llm}
+embabel.agent.rag.compression-llm.model=${guide.rag-llm}
+embabel.agent.rag.reranking-llm.model=${guide.rag-llm}
 
 # This is actually the default but it's good to show
 # How to change it
-embabel.models.default-embedding-model=nomic-embed-text:latest
+#embabel.models.default-embedding-model=nomic-embed-text:latest
+embabel.models.default-embedding-model=text-embedding-3-small
+
 
-guide.default-llm=qwen3:0.6b
+guide.rag-llm=gpt-4.1-nano
 
 guide.persona=bible
 guide.top-k=10
-guide.similarity-threshold=0.2
-guide.llm.model=qwen3:4b
+guide.similarity-threshold=0.7
+guide.llm.model=gpt-4.1-mini
+#guide.llm.model=qwen3:4b
+
 
 embabel.shell.line-length=140
 
-logging.level.com.embabel.agent.rag=INFO
+logging.level.com.embabel.agent.rag=INFO
+
+logging.level.com.embabel.agent.rag.tools=DEBUG
+
+logging.level.com.embabel.agent.rag.pipeline=DEBUG

src/main/resources/prompts/coding_style.jinja

@@ -1,7 +1,8 @@
 ## Coding Style
 
 Refer to this if you generate any code.
-TRY NOT TO GENERATE CODE BUT PULL FROM EXAMPLES.
+TRY NOT TO GENERATE CODE BUT PULL FROM EXAMPLES OR/AND DOCUMENTATION.
+CODE GENERATION IS A LAST RESORT.
 
 - Use modern Java features like var, records, and switch expressions.
 - Use multiline strings rather than concatenation.