Draft: RAG and kNN

assets/js/blocks/rag/block.json

@@ -0,0 +1,38 @@
+{
+	"$schema": "https://schemas.wp.org/trunk/block.json",
+	"apiVersion": 3,
+	"name": "elasticpress-labs/rag",
+	"title": "RAG",
+	"category": "elasticpress",
+	"icon": "superhero-alt",
+	"description": "RAG block description.",
+	"textdomain": "elasticpress-labs",
+	"attributes": {
+		"title": {
+			"default": "Title",
+			"type": "string"
+		}
+	},
+	"supports": {
+		"align": true,
+		"color": {
+			"background": true,
+			"link": true,
+			"text": false
+		},
+		"html": false,
+		"position": {
+			"sticky": true
+		},
+		"spacing": {
+			"margin": true,
+			"padding": true
+		},
+		"typography": {
+			"fontSize": true,
+			"lineHeight": true
+		}
+	},
+	"editorScript": "ep-rag-block-script",
+	"script": "ep-rag-block-frontend-script"
+}

assets/js/blocks/rag/edit.js

@@ -0,0 +1,33 @@
+/**
+ * WordPress dependencies.
+ */
+import { RichText, useBlockProps } from '@wordpress/block-editor';
+import { __ } from '@wordpress/i18n';
+
+/**
+ * Edit component.
+ *
+ * @param {object} props Component props.
+ * @param {object} props.attributes Block attributes.
+ * @param {Function} props.setAttributes Block attribute setter.
+ * @returns {Function} Component.
+ */
+export default ({ attributes, setAttributes }) => {
+	const { title } = attributes;
+
+	const blockProps = useBlockProps({
+		className: 'ep-rag',
+	});
+
+	return (
+		<div {...blockProps}>
+			<RichText
+				aria-label={__('Title text', 'elasticpress-labs')}
+				placeholder={__('Add title', 'elasticpress-labs')}
+				withoutInteractiveFormatting
+				value={title}
+				onChange={(html) => setAttributes({ title: html })}
+			/>
+		</div>
+	);
+};

assets/js/blocks/rag/frontend.js

@@ -0,0 +1,56 @@
+/**
+ * WordPress dependencies.
+ */
+import apiFetch from '@wordpress/api-fetch';
+import domReady from '@wordpress/dom-ready';
+import { Placeholder } from '@wordpress/components';
+import { createRoot, render, useEffect, useState, WPElement } from '@wordpress/element';
+import Skeleton from 'react-loading-skeleton';
+import 'react-loading-skeleton/dist/skeleton.css';
+
+const { searchQuery, restApiEndpoint } = window.epRag;
+
+/**
+ * App component
+ *
+ * @returns {WPElement} App component.
+ */
+const App = () => {
+	const [message, setMessage] = useState('');
+	const [isLoading, setIsLoading] = useState(true);
+
+	useEffect(() => {
+		apiFetch({
+			path: `${restApiEndpoint}?search_query=${searchQuery}`,
+		})
+			.then((response) => {
+				setMessage(response.html);
+			})
+			.finally(() => {
+				setIsLoading(false);
+			});
+	}, []);
+
+	return isLoading ? (
+		<Placeholder>
+			<Skeleton count={5} />
+		</Placeholder>
+	) : (
+		// eslint-disable-next-line react/no-danger
+		<p dangerouslySetInnerHTML={{ __html: message }} />
+	);
+};
+
+domReady(() => {
+	const ragBlocks = document.querySelectorAll('.ep-rag-response');
+
+	ragBlocks.forEach((ragBlock) => {
+		if (typeof createRoot === 'function') {
+			const root = createRoot(ragBlock);
+
+			root.render(<App />);
+		} else {
+			render(<App />, ragBlock);
+		}
+	});
+});

assets/js/blocks/rag/index.js

@@ -0,0 +1,15 @@
+/**
+ * WordPress dependencies.
+ */
+import { registerBlockType } from '@wordpress/blocks';
+
+/**
+ * Internal dependencies.
+ */
+import { name } from './block.json';
+import Edit from './edit';
+
+registerBlockType(name, {
+	edit: (props) => <Edit {...props} />,
+	save: () => {},
+});

includes/classes/Feature/KnnSearch/KnnSearch.php

@@ -0,0 +1,46 @@
+<?php
+/**
+ * kNN Search Feature
+ *
+ * @since 2.4.0
+ * @package ElasticPressLabs
+ */
+
+namespace ElasticPressLabs\Feature\KnnSearch;
+
+use ElasticPress\Feature;
+
+if ( ! defined( 'ABSPATH' ) ) {
+	exit; // Exit if accessed directly.
+}
+
+/**
+ * kNN Search Feature
+ */
+class KnnSearch extends Feature {
+	/**
+	 * Initialize feature setting it's config
+	 */
+	public function __construct() {
+		$this->slug = 'knn_search';
+
+		$this->title = esc_html__( 'kNN Search', 'elasticpress-labs' );
+
+		$this->summary = __( 'kNN Search.', 'elasticpress-labs' );
+
+		$this->requires_feature = 'vector_embeddings';
+
+		parent::__construct();
+	}
+
+	/**
+	 * Connects the Module with WordPress using Hooks and/or Filters.
+	 *
+	 * @return void
+	 */
+	public function setup() {
+		\ElasticPress\SearchAlgorithms::factory()->register( new SearchAlgorithm\Knn() );
+		\ElasticPress\SearchAlgorithms::factory()->register( new SearchAlgorithm\KnnCosine() );
+		\ElasticPress\SearchAlgorithms::factory()->register( new SearchAlgorithm\Hybrid() );
+	}
+}

includes/classes/Feature/KnnSearch/SearchAlgorithm/Hybrid.php

@@ -0,0 +1,79 @@
+<?php
+/**
+ * Hybrid search algorithm
+ *
+ * @since 2.4.0
+ * @package elasticpress
+ */
+
+namespace ElasticPressLabs\Feature\KnnSearch\SearchAlgorithm;
+
+if ( ! defined( 'ABSPATH' ) ) {
+	// @codeCoverageIgnoreStart
+	exit; // Exit if accessed directly.
+	// @codeCoverageIgnoreEnd
+}
+
+/**
+ * Hybrid search algorithm class.
+ */
+class Hybrid extends SearchAlgorithm {
+	/**
+	 * Search algorithm slug.
+	 *
+	 * @return string
+	 */
+	public function get_slug(): string {
+		return 'hybrid_knn';
+	}
+
+	/**
+	 * Search algorithm name.
+	 *
+	 * @return string
+	 */
+	public function get_name(): string {
+		return esc_html__( 'Hybrid (kNN + Regular ES)', 'elasticpress-labs' );
+	}
+
+	/**
+	 * Search algorithm description.
+	 *
+	 * @return string
+	 */
+	public function get_description(): string {
+		return esc_html__( 'Search using a mix of Elasticsearch kNN and a regular query.', 'elasticpress-labs' );
+	}
+
+	/**
+	 * Return the whole ES query
+	 *
+	 * @param array     $formatted_args Formatted Elasticsearch query
+	 * @param array     $args           The WP_Query variables
+	 * @param \WP_Query $query          The WP_Query object
+	 * @return array
+	 */
+	public function get_es_query( $formatted_args, $args, $query ): array {
+		if ( ! $query->is_search() ) {
+			return $formatted_args;
+		}
+
+		$query_embedding = $this->get_search_term_vector( $query->query_vars['s'] );
+		if ( empty( $query_embedding ) ) {
+			return $formatted_args;
+		}
+
+		return [
+			'from'        => $formatted_args['from'],
+			'size'        => $formatted_args['size'],
+			'post_filter' => $formatted_args['post_filter'],
+			'query'       => $formatted_args['query'],
+			'knn'         => [
+				'field'          => 'chunks.vector',
+				'query_vector'   => array_map( 'floatval', $query_embedding ),
+				'num_candidates' => 200,
+				'k'              => (int) $formatted_args['size'],
+			],
+		];
+	}
+}

includes/classes/Feature/KnnSearch/SearchAlgorithm/Knn.php

@@ -0,0 +1,78 @@
+<?php
+/**
+ * kNN search algorithm
+ *
+ * @since 2.4.0
+ * @package elasticpress
+ */
+
+namespace ElasticPressLabs\Feature\KnnSearch\SearchAlgorithm;
+
+if ( ! defined( 'ABSPATH' ) ) {
+	// @codeCoverageIgnoreStart
+	exit; // Exit if accessed directly.
+	// @codeCoverageIgnoreEnd
+}
+
+/**
+ * kNN search algorithm class.
+ */
+class Knn extends SearchAlgorithm {
+	/**
+	 * Search algorithm slug.
+	 *
+	 * @return string
+	 */
+	public function get_slug(): string {
+		return 'knn';
+	}
+
+	/**
+	 * Search algorithm name.
+	 *
+	 * @return string
+	 */
+	public function get_name(): string {
+		return esc_html__( 'kNN', 'elasticpress-labs' );
+	}
+
+	/**
+	 * Search algorithm description.
+	 *
+	 * @return string
+	 */
+	public function get_description(): string {
+		return esc_html__( 'Search using Elasticsearch kNN.', 'elasticpress-labs' );
+	}
+
+	/**
+	 * Return the whole ES query
+	 *
+	 * @param array     $formatted_args Formatted Elasticsearch query
+	 * @param array     $args           The WP_Query variables
+	 * @param \WP_Query $query          The WP_Query object
+	 * @return array
+	 */
+	public function get_es_query( $formatted_args, $args, $query ): array {
+		if ( ! $query->is_search() ) {
+			return $formatted_args;
+		}
+
+		$query_embedding = $this->get_search_term_vector( $query->query_vars['s'] );
+		if ( empty( $query_embedding ) ) {
+			return $formatted_args;
+		}
+
+		return [
+			'from'        => $formatted_args['from'],
+			'size'        => $formatted_args['size'],
+			'post_filter' => $formatted_args['post_filter'],
+			'knn'         => [
+				'field'          => 'chunks.vector',
+				'query_vector'   => array_map( 'floatval', $query_embedding ),
+				'num_candidates' => 200,
+				'k'              => (int) $formatted_args['size'],
+			],
+		];
+	}
+}