Merge pull request #22 from schemacrawler/full-text

Full text search for models that do not support embedding

Author: sualeh

pom.xml

@@ -81,6 +81,12 @@
       <artifactId>jackson-module-jsonSchema</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>us.fatehi</groupId>
+      <artifactId>full-text-search</artifactId>
+      <version>0.0.6</version>
+    </dependency>
+    
     <dependency>
       <groupId>org.apache.commons</groupId>
       <artifactId>commons-math3</artifactId>

src/main/assembly/assembly.xml

@@ -18,13 +18,20 @@
       <outputDirectory>.</outputDirectory>
       <useProjectArtifact>false</useProjectArtifact>
       <unpack>false</unpack>
+	  <useTransitiveFiltering>false</useTransitiveFiltering>
       <scope>runtime</scope>
       <excludes>
         <exclude>org.slf4j:*</exclude>
-        <exclude>commons-dbutils:*</exclude>
-		<exclude>com.fasterxml.jackson.core:*</exclude>
+		<exclude>org.junit.platform:*</exclude>
+		<exclude>org.junit.jupiter:*</exclude>
+		<exclude>org.opentest4j:*</exclude>
+		<exclude>com.fasterxml.jackson.*:*</exclude>
+		<exclude>org.yaml:*</exclude>
+		<exclude>javax.validation:*</exclude>
+		<exclude>com.azure:azure-core-test</exclude>
+		<exclude>org.apache.ant:*</exclude>
         <exclude>us.fatehi:*</exclude>
       </excludes>
-    </dependencySet>
+    </dependencySet>	
   </dependencySets>
 </assembly>

src/main/java/schemacrawler/tools/command/aichat/utility/lanchain4j/AiModelFactoryUtility.java

@@ -40,6 +40,8 @@ public class AiModelFactoryUtility {
 
   interface AiModelFactory {
 
+    boolean hasEmbeddingModel();
+
     boolean isSupported();
 
     ChatLanguageModel newChatLanguageModel();

src/main/java/schemacrawler/tools/command/aichat/utility/lanchain4j/AnthropicModelFactory.java

@@ -29,19 +29,14 @@
 package schemacrawler.tools.command.aichat.utility.lanchain4j;
 
 import java.time.Duration;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
 import static java.util.Objects.requireNonNull;
-import dev.langchain4j.data.embedding.Embedding;
-import dev.langchain4j.data.segment.TextSegment;
+import dev.langchain4j.exception.UnsupportedFeatureException;
 import dev.langchain4j.memory.ChatMemory;
 import dev.langchain4j.memory.chat.MessageWindowChatMemory;
 import dev.langchain4j.model.anthropic.AnthropicChatModel;
 import dev.langchain4j.model.anthropic.AnthropicChatModelName;
 import dev.langchain4j.model.chat.ChatLanguageModel;
 import dev.langchain4j.model.embedding.EmbeddingModel;
-import dev.langchain4j.model.output.Response;
 import schemacrawler.tools.command.aichat.options.AiChatCommandOptions;
 import schemacrawler.tools.command.aichat.utility.lanchain4j.AiModelFactoryUtility.AiModelFactory;
 import us.fatehi.utility.property.PropertyName;
@@ -55,6 +50,11 @@ public AnthropicModelFactory(final AiChatCommandOptions commandOptions) {
     aiChatCommandOptions = requireNonNull(commandOptions, "No AI Chat options provided");
   }
 
+  @Override
+  public boolean hasEmbeddingModel() {
+    return false;
+  }
+
   @Override
   public boolean isSupported() {
     if (!aiChatCommandOptions.aiProvider().equals(aiProvider.getName())) {
@@ -88,20 +88,7 @@ public ChatMemory newChatMemory() {
 
   @Override
   public EmbeddingModel newEmbeddingModel() {
-    return new EmbeddingModel() {
-
-      @Override
-      public Response<List<Embedding>> embedAll(final List<TextSegment> textSegments) {
-        if (textSegments == null || textSegments.isEmpty()) {
-          return new Response<>(Collections.emptyList());
-        }
-
-        final Embedding embedding = new Embedding(new float[] {0f});
-        final Embedding[] embeddings = new Embedding[textSegments.size()];
-        Arrays.fill(embeddings, embedding);
-        return new Response<>(Arrays.asList(embeddings));
-      }
-    };
+    throw new UnsupportedFeatureException("Anthropic does not have embedding models");
   }
 
   @Override

src/main/java/schemacrawler/tools/command/aichat/utility/lanchain4j/FullTextCatalogContentRetriever.java

@@ -0,0 +1,64 @@
+package schemacrawler.tools.command.aichat.utility.lanchain4j;
+
+import java.util.List;
+import java.util.logging.Logger;
+import org.apache.lucene.store.Directory;
+import static java.util.Objects.requireNonNull;
+import dev.langchain4j.data.document.Metadata;
+import dev.langchain4j.data.segment.TextSegment;
+import dev.langchain4j.rag.content.Content;
+import dev.langchain4j.rag.content.retriever.ContentRetriever;
+import dev.langchain4j.rag.query.Query;
+import schemacrawler.schema.Catalog;
+import schemacrawler.schema.DatabaseInfo;
+import schemacrawler.schema.Table;
+import schemacrawler.tools.command.serialize.model.CompactCatalogUtility;
+import schemacrawler.tools.command.serialize.model.TableDocument;
+import us.fatehi.search.DirectoryFactory;
+import us.fatehi.search.LuceneContentRetriever;
+import us.fatehi.search.LuceneIndexer;
+
+public class FullTextCatalogContentRetriever implements ContentRetriever {
+
+  private static final Logger LOGGER =
+      Logger.getLogger(Langchain4JChatAssistant.class.getCanonicalName());
+
+  private final LuceneContentRetriever fullTextCatalogRetriever;
+
+  public FullTextCatalogContentRetriever(final Catalog catalog) {
+    requireNonNull(catalog, "No catalog provided");
+
+    final Directory tempDirectory = DirectoryFactory.tempDirectory();
+    final LuceneIndexer luceneIndexer = new LuceneIndexer(tempDirectory);
+    final TextSegment databaseInfoContent = getDatabaseInfoContent(catalog);
+    luceneIndexer.addContent(databaseInfoContent);
+    for (final Table table : catalog.getTables()) {
+      final TableDocument tableDocument = CompactCatalogUtility.getTableDocument(table, false);
+      luceneIndexer.addContent(TextSegment.from(tableDocument.toJson()));
+    }
+
+    fullTextCatalogRetriever =
+        LuceneContentRetriever.builder()
+            .directory(tempDirectory)
+            .matchUntilTopN()
+            .maxTokens(5_000)
+            .build();
+  }
+
+  @Override
+  public List<Content> retrieve(final Query query) {
+    return fullTextCatalogRetriever.retrieve(query);
+  }
+
+  private TextSegment getDatabaseInfoContent(final Catalog catalog) {
+    final DatabaseInfo databaseInfo = catalog.getDatabaseInfo();
+    final String databaseProductName = databaseInfo.getDatabaseProductName();
+    final Metadata metadata = new Metadata();
+    metadata.put("database", databaseProductName);
+    metadata.put("database-version", databaseInfo.getDatabaseProductVersion());
+    final TextSegment textSegment =
+        TextSegment.from(
+            String.format("Customize SQL queries for %s", databaseProductName), metadata);
+    return textSegment;
+  }
+}

src/main/java/schemacrawler/tools/command/aichat/utility/lanchain4j/GitHubModelFactory.java

@@ -52,6 +52,11 @@ public GitHubModelFactory(final AiChatCommandOptions commandOptions) {
     aiChatCommandOptions = requireNonNull(commandOptions, "No AI Chat options provided");
   }
 
+  @Override
+  public boolean hasEmbeddingModel() {
+    return true;
+  }
+
   @Override
   public boolean isSupported() {
     if (!aiChatCommandOptions.aiProvider().equals(aiProvider.getName())) {

src/main/java/schemacrawler/tools/command/aichat/utility/lanchain4j/Langchain4JChatAssistant.java

@@ -97,8 +97,12 @@ public Langchain4JChatAssistant(
 
     final boolean useMetadata = aiChatOptions.useMetadata();
     if (useMetadata) {
-      final EmbeddingModel embeddingModel = modelFactory.newEmbeddingModel();
-      contentRetriever = new CatalogContentRetriever(embeddingModel, catalog);
+      if (modelFactory.hasEmbeddingModel()) {
+        final EmbeddingModel embeddingModel = modelFactory.newEmbeddingModel();
+        contentRetriever = new CatalogContentRetriever(embeddingModel, catalog);
+      } else {
+        contentRetriever = new FullTextCatalogContentRetriever(catalog);
+      }
     } else {
       contentRetriever = query -> Collections.emptyList();
     }
@@ -171,7 +175,7 @@ private SystemMessage createSystemMessage(final String prompt) {
     buffer.append(metadataPriming).append("\n");
     final List<Content> contents = contentRetriever.retrieve(Query.from(prompt));
     for (final Content content : contents) {
-      buffer.append(content.textSegment().text()).append("\n");
+      buffer.append("\n").append(content.textSegment().text());
     }
 
     return SystemMessage.from(buffer.toString());

src/main/java/schemacrawler/tools/command/aichat/utility/lanchain4j/OpenAIModelFactory.java

@@ -53,6 +53,11 @@ public OpenAIModelFactory(final AiChatCommandOptions commandOptions) {
     aiChatCommandOptions = requireNonNull(commandOptions, "No AI Chat options provided");
   }
 
+  @Override
+  public boolean hasEmbeddingModel() {
+    return true;
+  }
+
   @Override
   public boolean isSupported() {
     if (!aiChatCommandOptions.aiProvider().equals(aiProvider.getName())) {