From 205ff82db9580f08c3b3ccdba80178c2d5e1c3fa Mon Sep 17 00:00:00 2001
From: Christopher Hertel <mail@christopher-hertel.de>
Date: Sat, 21 Sep 2024 23:57:26 +0200
Subject: [PATCH] feat: adding basic similarity search as tool

---
 README.md                             |  3 +-
 src/Store/Azure/SearchStore.php       |  2 +-
 src/Store/ChromaDb/Store.php          |  2 +-
 src/Store/VectorStoreInterface.php    |  6 ++--
 src/ToolBox/Tool/SimilaritySearch.php | 45 +++++++++++++++++++++++++++
 5 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 src/ToolBox/Tool/SimilaritySearch.php

diff --git a/README.md b/README.md
index c9b101a..48d888f 100644
--- a/README.md
+++ b/README.md
@@ -50,8 +50,9 @@ Currently supported stores:
 Provided Tools
 --------------
 
-* [x] SerpApi
 * [x] Clock
+* [x] SerpApi
+* [x] Similarity Search (Basic)
 * [x] Wikipedia
 * [x] Weather
 * [x] YouTube Transcriber
diff --git a/src/Store/Azure/SearchStore.php b/src/Store/Azure/SearchStore.php
index 9cf6a50..2a727c5 100644
--- a/src/Store/Azure/SearchStore.php
+++ b/src/Store/Azure/SearchStore.php
@@ -37,7 +37,7 @@ public function addDocuments(array $documents): void
     /**
      * @return list<Document>
      */
-    public function query(Vector $vector): array
+    public function query(Vector $vector, array $options = []): array
     {
         $result = $this->request('search', [
             'vectorQueries' => [$this->buildVectorQuery($vector)],
diff --git a/src/Store/ChromaDb/Store.php b/src/Store/ChromaDb/Store.php
index 1e2a9a2..106efed 100644
--- a/src/Store/ChromaDb/Store.php
+++ b/src/Store/ChromaDb/Store.php
@@ -45,7 +45,7 @@ public function addDocuments(array $documents): void
         $collection->add($ids, $vectors, $metadata);
     }
 
-    public function query(Vector $vector): array
+    public function query(Vector $vector, array $options = []): array
     {
         $collection = $this->client->getOrCreateCollection($this->collectionName);
         $queryResponse = $collection->query(
diff --git a/src/Store/VectorStoreInterface.php b/src/Store/VectorStoreInterface.php
index 5f9e7d6..0843331 100644
--- a/src/Store/VectorStoreInterface.php
+++ b/src/Store/VectorStoreInterface.php
@@ -10,7 +10,9 @@
 interface VectorStoreInterface extends StoreInterface
 {
     /**
-     * @return list<Document>
+     * @param array<string, mixed> $options
+     *
+     * @return Document[]
      */
-    public function query(Vector $vector): array;
+    public function query(Vector $vector, array $options = []): array;
 }
diff --git a/src/ToolBox/Tool/SimilaritySearch.php b/src/ToolBox/Tool/SimilaritySearch.php
new file mode 100644
index 0000000..fe368cf
--- /dev/null
+++ b/src/ToolBox/Tool/SimilaritySearch.php
@@ -0,0 +1,45 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\ToolBox\Tool;
+
+use PhpLlm\LlmChain\Document\Document;
+use PhpLlm\LlmChain\EmbeddingModel;
+use PhpLlm\LlmChain\Store\VectorStoreInterface;
+use PhpLlm\LlmChain\ToolBox\AsTool;
+
+#[AsTool('similarity_search', description: 'Searches for documents similar to a query or sentence.')]
+final class SimilaritySearch
+{
+    /**
+     * @var Document[]
+     */
+    public array $usedDocuments = [];
+
+    public function __construct(
+        private readonly EmbeddingModel $embedding,
+        private readonly VectorStoreInterface $vectorStore,
+    ) {
+    }
+
+    /**
+     * @param string $searchTerm string used for similarity search
+     */
+    public function __invoke(string $searchTerm): string
+    {
+        $vector = $this->embedding->create($searchTerm);
+        $this->usedDocuments = $this->vectorStore->query($vector);
+
+        if (0 === count($this->usedDocuments)) {
+            return 'No results found';
+        }
+
+        $result = 'Found documents with following information:'.PHP_EOL;
+        foreach ($this->usedDocuments as $document) {
+            $result .= json_encode($document->metadata);
+        }
+
+        return $result;
+    }
+}