[AUR-389] Add base interface and embedding model (#17)

This change provides the base interface of an embedding, and wrap the Langchain's OpenAI embedding. Usage as follow:

```python
from kotaemon.embeddings import AzureOpenAIEmbeddings

model = AzureOpenAIEmbeddings(
    model="text-embedding-ada-002",
    deployment="embedding-deployment",
    openai_api_base="https://test.openai.azure.com/",
    openai_api_key="some-key",
)
output = model("Hello world")
```

knowledgehub/embeddings/base.py

@@ -0,0 +1,62 @@
+from typing import List, Type
+
+from langchain.embeddings.base import Embeddings as LCEmbeddings
+from theflow import Param
+
+from ..components import BaseComponent
+from ..documents.base import Document
+
+
+class Embeddings(BaseComponent):
+    ...
+
+
+class LangchainEmbeddings(Embeddings):
+    _lc_class: Type[LCEmbeddings]
+
+    def __init__(self, **params):
+        if self._lc_class is None:
+            raise AttributeError(
+                "Should set _lc_class attribute to the LLM class from Langchain "
+                "if using LLM from Langchain"
+            )
+
+        self._kwargs: dict = {}
+        for param in list(params.keys()):
+            if param in self._lc_class.__fields__:  # type: ignore
+                self._kwargs[param] = params.pop(param)
+        super().__init__(**params)
+
+    def __setattr__(self, name, value):
+        if name in self._lc_class.__fields__:
+            setattr(self.agent, name, value)
+        else:
+            super().__setattr__(name, value)
+
+    @Param.decorate(no_cache=True)
+    def agent(self):
+        return self._lc_class(**self._kwargs)
+
+    def run_raw(self, text: str) -> List[float]:
+        return self.agent.embed_query(text)  # type: ignore
+
+    def run_batch_raw(self, text: List[str]) -> List[List[float]]:
+        return self.agent.embed_documents(text)  # type: ignore
+
+    def run_document(self, text: Document) -> List[float]:
+        return self.agent.embed_query(text.text)  # type: ignore
+
+    def run_batch_document(self, text: List[Document]):
+        return self.agent.embed_documents([each.text for each in text])  # type: ignore
+
+    def is_document(self, text) -> bool:
+        if isinstance(text, Document):
+            return True
+        elif isinstance(text, List) and isinstance(text[0], Document):
+            return True
+        return False
+
+    def is_batch(self, text) -> bool:
+        if isinstance(text, list):
+            return True
+        return False

knowledgehub/embeddings/openai.py

@@ -0,0 +1,15 @@
+from langchain.embeddings import OpenAIEmbeddings as LCOpenAIEmbeddings
+
+from .base import LangchainEmbeddings
+
+
+class OpenAIEmbeddings(LangchainEmbeddings):
+    _lc_class = LCOpenAIEmbeddings
+
+
+class AzureOpenAIEmbeddings(LangchainEmbeddings):
+    _lc_class = LCOpenAIEmbeddings
+
+    def __init__(self, **params):
+        params["openai_api_type"] = "azure"
+        super().__init__(**params)

knowledgehub/llms/chats/base.py

@@ -30,8 +30,8 @@ class LangchainChatLLM(ChatLLM):
                 self._kwargs[param] = params.pop(param)
         super().__init__(**params)
 
-    @Param.decorate()
-    def agent(self):
+    @Param.decorate(no_cache=True)
+    def agent(self) -> BaseLanguageModel:
         return self._lc_class(**self._kwargs)
 
     def run_raw(self, text: str) -> LLMInterface:
@@ -43,7 +43,7 @@ class LangchainChatLLM(ChatLLM):
         return self.run_batch_document(inputs)
 
     def run_document(self, text: List[Message]) -> LLMInterface:
-        pred = self.agent.generate([text])
+        pred = self.agent.generate([text])  # type: ignore
         return LLMInterface(
             text=[each.text for each in pred.generations[0]],
             completion_tokens=pred.llm_output["token_usage"]["completion_tokens"],

tests/test_embedding_models.py

@@ -0,0 +1,46 @@
+import json
+from pathlib import Path
+from unittest.mock import patch
+
+from kotaemon.embeddings.openai import AzureOpenAIEmbeddings
+
+with open(Path(__file__).parent / "resources" / "embedding_openai_batch.json") as f:
+    openai_embedding_batch = json.load(f)
+
+with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
+    openai_embedding = json.load(f)
+
+
+@patch(
+    "openai.api_resources.embedding.Embedding.create",
+    side_effect=lambda *args, **kwargs: openai_embedding,
+)
+def test_azureopenai_embeddings_raw(openai_embedding_call):
+    model = AzureOpenAIEmbeddings(
+        model="text-embedding-ada-002",
+        deployment="embedding-deployment",
+        openai_api_base="https://test.openai.azure.com/",
+        openai_api_key="some-key",
+    )
+    output = model("Hello world")
+    assert isinstance(output, list)
+    assert isinstance(output[0], float)
+    openai_embedding_call.assert_called()
+
+
+@patch(
+    "openai.api_resources.embedding.Embedding.create",
+    side_effect=lambda *args, **kwargs: openai_embedding_batch,
+)
+def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
+    model = AzureOpenAIEmbeddings(
+        model="text-embedding-ada-002",
+        deployment="embedding-deployment",
+        openai_api_base="https://test.openai.azure.com/",
+        openai_api_key="some-key",
+    )
+    output = model(["Hello world", "Goodbye world"])
+    assert isinstance(output, list)
+    assert isinstance(output[0], list)
+    assert isinstance(output[0][0], float)
+    openai_embedding_call.assert_called()