feat(theodo-group#111): add dimensions to interface

bernard-ng · May 17, 2024 · a54da80 · a54da80
1 parent e0f1425
commit a54da80
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 13 deletions.
diff --git a/src/Embeddings/EmbeddingGenerator/EmbeddingGeneratorInterface.php b/src/Embeddings/EmbeddingGenerator/EmbeddingGeneratorInterface.php
@@ -9,15 +9,15 @@ interface EmbeddingGeneratorInterface
  /**
  * @return float[]
  */
- public function embedText(string $text): array;
+ public function embedText(string $text, ?int $dimensions = null): array;
 
- public function embedDocument(Document $document): Document;
+ public function embedDocument(Document $document, ?int $dimensions = null): Document;
 
  /**
  * @param Document[] $documents
  * @return Document[]
  */
- public function embedDocuments(array $documents): array;
+ public function embedDocuments(array $documents, ?int $dimensions = null): array;
 
  public function getEmbeddingLength(): int;
 }
diff --git a/src/Embeddings/EmbeddingGenerator/Mistral/MistralEmbeddingGenerator.php b/src/Embeddings/EmbeddingGenerator/Mistral/MistralEmbeddingGenerator.php
@@ -37,7 +37,7 @@ public function __construct(?OpenAIConfig $config = null)
  *
  * @return float[]
  */
- public function embedText(string $text): array
+ public function embedText(string $text, ?int $dimensions = null): array
  {
  $text = str_replace("\n", ' ', $text);
 
@@ -64,10 +64,10 @@ public function embedText(string $text): array
  return $searchResults['data'][0]['embedding'];
  }
 
- public function embedDocument(Document $document): Document
+ public function embedDocument(Document $document, ?int $dimensions = null): Document
  {
  $text = $document->formattedContent ?? $document->content;
- $document->embedding = $this->embedText($text);
+ $document->embedding = $this->embedText($text, $dimensions);
 
  return $document;
  }
@@ -78,11 +78,11 @@ public function embedDocument(Document $document): Document
  * @param Document[] $documents
  * @return Document[]
  */
- public function embedDocuments(array $documents): array
+ public function embedDocuments(array $documents, ?int $dimensions = null): array
  {
  $embedDocuments = [];
  foreach ($documents as $document) {
- $embedDocuments[] = $this->embedDocument($document);
+ $embedDocuments[] = $this->embedDocument($document, $dimensions);
  }
 
  return $embedDocuments;

diff --git a/src/Embeddings/EmbeddingGenerator/Ollama/OllamaEmbeddingGenerator.php b/src/Embeddings/EmbeddingGenerator/Ollama/OllamaEmbeddingGenerator.php
@@ -30,8 +30,11 @@ public function __construct(OllamaConfig $config)
  * Call out to Ollama embedding endpoint.
  *
  * @return float[]
+ *
+ * @throws \JsonException
+ * @throws Exception
  */
- public function embedText(string $text): array
+ public function embedText(string $text, ?int $dimensions = null): array
  {
  $text = str_replace("\n", ' ', $text);
 
@@ -57,23 +60,28 @@ public function embedText(string $text): array
  return $searchResults['embedding'];
  }
 
- public function embedDocument(Document $document): Document
+ /**
+ * @throws \JsonException
+ */
+ public function embedDocument(Document $document, ?int $dimensions = null): Document
  {
  $text = $document->formattedContent ?? $document->content;
- $document->embedding = $this->embedText($text);
+ $document->embedding = $this->embedText($text, $dimensions);
 
  return $document;
  }
 
  /**
  * @param Document[] $documents
  * @return Document[]
+ *
+ * @throws \JsonException
  */
- public function embedDocuments(array $documents): array
+ public function embedDocuments(array $documents, ?int $dimensions = null): array
  {
  $embedDocuments = [];
  foreach ($documents as $document) {
- $embedDocuments[] = $this->embedDocument($document);
+ $embedDocuments[] = $this->embedDocument($document, $dimensions);
  }
 
  return $embedDocuments;