FakeEmbeddingModel API

import com.embabel.common.test.ai.FakeEmbeddingModel

import com.embabel.common.test.ai.FakeEmbeddingModel;

/**
 * Fake embedding model that implements Spring AI's EmbeddingModel interface
 * for testing purposes. Generates random embeddings without making actual API calls.
 *
 * @property dimensions The dimension size of generated embedding vectors (default: 1536)
 */
data class FakeEmbeddingModel(
    val dimensions: Int = 1536
) : EmbeddingModel

// Default dimensions (1536)
val model = FakeEmbeddingModel()

// Custom dimensions
val model384 = FakeEmbeddingModel(dimensions = 384)
val model768 = FakeEmbeddingModel(dimensions = 768)

override fun embed(document: Document): FloatArray

val model = FakeEmbeddingModel(dimensions = 768)
val document = Document("Sample text for embedding")
val embedding = model.embed(document)
// Returns a random FloatArray of size 768

FakeEmbeddingModel model = new FakeEmbeddingModel(768);
Document document = new Document("Sample text for embedding");
float[] embedding = model.embed(document);
// Returns a random float[] of size 768

override fun embed(texts: List<String>): MutableList<FloatArray>

val model = FakeEmbeddingModel()
val texts = listOf(
    "First document",
    "Second document",
    "Third document"
)
val embeddings = model.embed(texts)
// Returns a list of 3 random embeddings, each of size 1536

FakeEmbeddingModel model = new FakeEmbeddingModel(1536);
List<String> texts = List.of(
    "First document",
    "Second document",
    "Third document"
);
List<float[]> embeddings = model.embed(texts);
// Returns a list of 3 random embeddings

override fun call(request: EmbeddingRequest): EmbeddingResponse

val model = FakeEmbeddingModel(dimensions = 384)
val request = EmbeddingRequest(
    listOf("query 1", "query 2"),
    null  // options
)
val response = model.call(request)

// Access embeddings from response
val embeddings = response.results
println("Generated ${embeddings.size} embeddings")
embeddings.forEach { embedding ->
    println("Embedding index: ${embedding.index}, dims: ${embedding.output.size}")
}

FakeEmbeddingModel model = new FakeEmbeddingModel(384);
EmbeddingRequest request = new EmbeddingRequest(
    List.of("query 1", "query 2"),
    null  // options
);
EmbeddingResponse response = model.call(request);

// Access embeddings from response
List<Embedding> embeddings = response.getResults();
System.out.println("Generated " + embeddings.size() + " embeddings");

@Test
fun `test basic embedding`() {
    val model = FakeEmbeddingModel(dimensions = 512)
    val document = Document("Test content")

    val embedding = model.embed(document)

    assertEquals(512, embedding.size)
    assertTrue(embedding.all { it.isFinite() })
}

@Test
fun `test batch embedding`() {
    val model = FakeEmbeddingModel()
    val texts = listOf("doc 1", "doc 2", "doc 3")

    val embeddings = model.embed(texts)

    assertEquals(3, embeddings.size)
    embeddings.forEach { embedding ->
        assertEquals(1536, embedding.size)
    }
}

@Test
fun `test embedding service`() {
    val embeddingModel = FakeEmbeddingModel(dimensions = 768)
    val service = MyEmbeddingService(embeddingModel)

    val result = service.processDocument(Document("test"))

    assertNotNull(result)
    assertEquals(768, result.embedding.size)
}

@Test
fun `test with embedding request`() {
    val model = FakeEmbeddingModel(dimensions = 384)
    val request = EmbeddingRequest(
        listOf("text 1", "text 2"),
        null
    )

    val response = model.call(request)

    assertEquals(2, response.results.size)
    response.results.forEach {
        assertEquals(384, it.output.size)
    }
}

@Test
fun `test with different dimension sizes`() {
    val model384 = FakeEmbeddingModel(dimensions = 384)
    val model768 = FakeEmbeddingModel(dimensions = 768)
    val model1536 = FakeEmbeddingModel(dimensions = 1536)

    val text = "sample text"
    val doc = Document(text)

    assertEquals(384, model384.embed(doc).size)
    assertEquals(768, model768.embed(doc).size)
    assertEquals(1536, model1536.embed(doc).size)
}

@TestConfiguration
class TestConfig {
    @Bean
    fun embeddingModel(): EmbeddingModel {
        return FakeEmbeddingModel(dimensions = 768)
    }
}

@SpringBootTest
class EmbeddingIntegrationTest {

    @Autowired
    private lateinit var embeddingModel: EmbeddingModel

    @Test
    fun `test with Spring bean`() {
        val embedding = embeddingModel.embed(Document("test"))
        assertNotNull(embedding)
    }
}

/**
 * Document representation in Spring AI
 */
class Document(
    val content: String,
    val metadata: Map<String, Any> = emptyMap()
)

/**
 * Request for embedding operations
 */
class EmbeddingRequest(
    val inputs: List<String>,
    val options: EmbeddingOptions?
)

/**
 * Response from embedding operations
 */
class EmbeddingResponse(
    val results: List<Embedding>
)

/**
 * Individual embedding result
 */
class Embedding(
    val output: FloatArray,
    val index: Int
)

/**
 * Spring AI's embedding model interface
 */
interface EmbeddingModel {
    fun embed(document: Document): FloatArray
    fun embed(texts: List<String>): MutableList<FloatArray>
    fun call(request: EmbeddingRequest): EmbeddingResponse
}