feat(llama-index): add embeddings

examples/extensions/llama_index/llama_index_embedding.py

@@ -0,0 +1,34 @@
+"""
+LlamaIndex Embeddings
+"""
+from dotenv import load_dotenv
+
+from genai import Client, Credentials
+from genai.extensions.llama_index import IBMGenAILlamaIndexEmbedding
+from genai.schema import TextEmbeddingParameters
+
+# make sure you have a .env file under genai root with
+# GENAI_KEY=<your-genai-key>
+# GENAI_API=<genai-api-endpoint> (optional) DEFAULT_API = "https://bam-api.res.ibm.com"
+load_dotenv()
+
+
+def heading(text: str) -> str:
+    """Helper function for centering text."""
+    return "\n" + f" {text} ".center(80, "=") + "\n"
+
+
+print(heading("LlamaIndex Embeddings"))
+
+client = Client(credentials=Credentials.from_env())
+embeddings = IBMGenAILlamaIndexEmbedding(
+    client=client,
+    model_id="sentence-transformers/all-minilm-l6-v2",
+    parameters=TextEmbeddingParameters(truncate_input_tokens=True),
+)
+
+query_embedding = embeddings.get_query_embedding("Hello world!")
+print(query_embedding)
+
+documents_embedding = embeddings.get_agg_embedding_from_queries(["First document", "Second document"])
+print(documents_embedding)

src/genai/extensions/llama_index/__init__.py

@@ -1,5 +1,6 @@
 """Extension for LLamaIndex library"""
 
+from genai.extensions.llama_index.embeddings import IBMGenAILlamaIndexEmbedding
 from genai.extensions.llama_index.llm import IBMGenAILlamaIndex
 
-__all__ = ["IBMGenAILlamaIndex"]
+__all__ = ["IBMGenAILlamaIndex", "IBMGenAILlamaIndexEmbedding"]

src/genai/extensions/llama_index/embeddings.py

@@ -0,0 +1,53 @@
+import asyncio
+from typing import Optional
+
+from genai._types import ModelLike
+from genai.client import Client
+from genai.schema import TextEmbeddingParameters
+from genai.text.embedding.embedding_service import CreateExecutionOptions
+
+try:
+    from llama_index.embeddings.base import BaseEmbedding, Embedding
+except ImportError:
+    raise ImportError("Could not import llamaindex: Please install ibm-generative-ai[llama-index] extension.")  # noqa: B904
+
+
+class IBMGenAILlamaIndexEmbedding(BaseEmbedding):
+    client: Client
+    model_id: str
+    parameters: Optional[ModelLike[TextEmbeddingParameters]] = None
+    execution_options: Optional[ModelLike[CreateExecutionOptions]] = None
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "IBMGenAIEmbedding"
+
+    def _get_query_embedding(self, query: str) -> Embedding:
+        response = self._get_embeddings([query])
+        return response[0]
+
+    def _get_text_embedding(self, text: str) -> Embedding:
+        response = self._get_embeddings([text])
+        return response[0]
+
+    def _get_text_embeddings(self, texts: list[str]) -> list[Embedding]:
+        response = self._get_embeddings(texts)
+        return response
+
+    async def _aget_query_embedding(self, query: str) -> Embedding:
+        return await asyncio.get_running_loop().run_in_executor(None, self._get_query_embedding, query)
+
+    async def _aget_text_embedding(self, text: str) -> Embedding:
+        return await asyncio.get_running_loop().run_in_executor(None, self._get_text_embedding, text)
+
+    async def _aget_text_embeddings(self, texts: list[str]) -> list[Embedding]:
+        return await asyncio.get_running_loop().run_in_executor(None, self._get_text_embeddings, texts)
+
+    def _get_embeddings(self, texts: list[str]) -> list[Embedding]:
+        embeddings: list[list[float]] = []
+        for response in self.client.text.embedding.create(
+            model_id=self.model_id, inputs=texts, parameters=self.parameters, execution_options=self.execution_options
+        ):
+            embeddings.extend(response.results)
+
+        return embeddings