Merge pull request #8 from CodeWithKyrian/ollama-embedding-function

Add support for Ollama Embedding Function

Author: CodeWithKyrian

.github/workflows/test.yml

@@ -16,12 +16,12 @@ jobs:
 
     services:
       chroma-wo-auth:
-        image: chromadb/chroma
+        image: chromadb/chroma:0.5.0
         ports:
           - 8000:8000
 
       chroma-w-auth:
-        image: chromadb/chroma
+        image: chromadb/chroma:0.5.0
         ports:
           - 8001:8000
         env:

docker-compose.yml

@@ -2,12 +2,12 @@ version: '3.9'
 
 services:
   chroma_wo_auth:
-    image: 'chromadb/chroma'
+    image: 'chromadb/chroma:0.5.0'
     ports:
       - '8000:8000'
 
   chroma_w_auth:
-    image: 'chromadb/chroma'
+    image: 'chromadb/chroma:0.5.0'
     ports:
       - '8001:8000'
     environment:

examples/index.php

@@ -6,6 +6,7 @@
 
 use Codewithkyrian\ChromaDB\ChromaDB;
 use Codewithkyrian\ChromaDB\Embeddings\JinaEmbeddingFunction;
+use Codewithkyrian\ChromaDB\Embeddings\OllamaEmbeddingFunction;
 
 $chroma = ChromaDB::factory()
     ->withDatabase('test_database')
@@ -14,22 +15,21 @@
 
 $chroma->deleteAllCollections();
 
-$embeddingFunction = new JinaEmbeddingFunction(
-    'jina_8cbaafb9543e42f1a2fc7430d456c3faKKPc93W8Eur5T2XjAkryfwQ9TOv8'
-);
+$embeddingFunction = new OllamaEmbeddingFunction();
 
 $collection = $chroma->createCollection(
     name: 'test_collection',
     embeddingFunction: $embeddingFunction
 );
 
+
 $collection->add(
-    ids: ['hello', 'world'],
-    documents: ['This is a test document', 'The man is happy']
+    ids: ['1', '2', '3'],
+    documents: ['He seems very happy', 'He was very sad when we last talked', 'She made him angry']
 );
 
 $queryResponse = $collection->query(
-    queryTexts: ['The man is excited'],
+    queryTexts: ['She annoyed him'],
     include: ['documents', 'distances']
 );
 

src/Embeddings/HuggingFaceEmbeddingServerFunction.php

@@ -28,9 +28,9 @@ public function generate(array $texts): array
 
         try {
             $response = $client->post('embed', [
-                'body' => json_encode([
+                'json' => [
                     'inputs' => $texts,
-                ])
+                ]
             ]);
         } catch (GuzzleException $e) {
             throw new \RuntimeException('Failed to generate embeddings', 0, $e);

src/Embeddings/OllamaEmbeddingFunction.php

@@ -0,0 +1,53 @@
+<?php
+
+declare(strict_types=1);
+
+
+namespace Codewithkyrian\ChromaDB\Embeddings;
+
+use GuzzleHttp\Client;
+
+class OllamaEmbeddingFunction implements EmbeddingFunction
+{
+    private Client $client;
+
+    public function __construct(
+        public readonly string $baseUrl = 'http://localhost:11434',
+        public readonly string $model = 'all-minilm',
+    )
+    {
+        $this->client = new Client([
+            'base_uri' => $this->baseUrl,
+            'headers' => [
+                'Content-Type' => 'application/json',
+            ]
+        ]);
+    }
+
+    /**
+     * @inheritDoc
+     */
+    public function generate(array $texts): array
+    {
+        try {
+            $embeddings = [];
+
+            foreach ($texts as $text) {
+                $response = $this->client->post('api/embeddings', [
+                    'json' => [
+                        'prompt' => $text,
+                        'model' => $this->model,
+                    ]
+                ]);
+
+                $result = json_decode($response->getBody()->getContents(), true);
+
+                $embeddings[] = $result['embedding'];
+            }
+
+            return $embeddings;
+        } catch (\Exception $e) {
+            throw new \RuntimeException('Failed to generate embeddings', 0, $e);
+        }
+    }
+}