cli[patch]: add embedding to integration template (#14881)

libs/cli/langchain_cli/integration_template/docs/text_embedding.ipynb

@@ -0,0 +1,132 @@
+{
+ "cells": [
+  {
+   "cell_type": "raw",
+   "id": "afaf8039",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "sidebar_label: __ModuleName__\n",
+    "---"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e49f1e0d",
+   "metadata": {},
+   "source": [
+    "# __ModuleName__Embeddings\n",
+    "\n",
+    "This notebook covers how to get started with __ModuleName__ embedding models.\n",
+    "\n",
+    "## Installation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c3bef91",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# install package\n",
+    "!pip install -U __package_name__"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2b4f3e15",
+   "metadata": {},
+   "source": [
+    "## Environment Setup\n",
+    "\n",
+    "Make sure to set the following environment variables:\n",
+    "\n",
+    "- TODO: fill out relevant environment variables or secrets\n",
+    "\n",
+    "## Usage"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62e0dbc3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from __module_name__.embeddings import __ModuleName__Embeddings\n",
+    "\n",
+    "embeddings = __ModuleName__Embeddings()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12fcfb4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "embeddings.embed_query(\"My query to look up\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1f2e6104",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "embeddings.embed_documents(\n",
+    "    [\"This is a content of the document\", \"This is another document\"]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46739f68",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# async embed query\n",
+    "await embeddings.aembed_query(\"My query to look up\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e48632ea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# async embed documents\n",
+    "await embeddings.aembed_documents(\n",
+    "    [\"This is a content of the document\", \"This is another document\"]\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

libs/cli/langchain_cli/integration_template/integration_template/__init__.py

@@ -1,5 +1,11 @@
 from __module_name__.chat_models import Chat__ModuleName__
+from __module_name__.embeddings import __ModuleName__Embeddings
 from __module_name__.llms import __ModuleName__LLM
 from __module_name__.vectorstores import __ModuleName__VectorStore
 
-__all__ = ["__ModuleName__LLM", "Chat__ModuleName__", "__ModuleName__VectorStore"]
+__all__ = [
+    "__ModuleName__LLM",
+    "Chat__ModuleName__",
+    "__ModuleName__VectorStore",
+    "__ModuleName__Embeddings",
+]

libs/cli/langchain_cli/integration_template/integration_template/embeddings.py

@@ -0,0 +1,34 @@
+from typing import List
+
+from langchain_core.embeddings import Embeddings
+
+
+class __ModuleName__Embeddings(Embeddings):
+    """__ModuleName__Embeddings embedding model.
+
+    Example:
+        .. code-block:: python
+
+            from __module_name__ import __ModuleName__Embeddings
+
+            model = __ModuleName__Embeddings()
+    """
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed search docs."""
+        raise NotImplementedError
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed query text."""
+        raise NotImplementedError
+
+    # only keep aembed_documents and aembed_query if they're implemented!
+    # delete them otherwise to use the base class' default
+    # implementation, which calls the sync version in an executor
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Asynchronous Embed search docs."""
+        raise NotImplementedError
+
+    async def aembed_query(self, text: str) -> List[float]:
+        """Asynchronous Embed query text."""
+        raise NotImplementedError

libs/cli/langchain_cli/integration_template/tests/integration_tests/test_embeddings.py

@@ -0,0 +1,19 @@
+"""Test __ModuleName__ embeddings."""
+from __module_name__.embeddings import __ModuleName__Embeddings
+
+
+def test___module_name___embedding_documents() -> None:
+    """Test cohere embeddings."""
+    documents = ["foo bar"]
+    embedding = __ModuleName__Embeddings()
+    output = embedding.embed_documents(documents)
+    assert len(output) == 1
+    assert len(output[0]) > 0
+
+
+def test___module_name___embedding_query() -> None:
+    """Test cohere embeddings."""
+    document = "foo bar"
+    embedding = __ModuleName__Embeddings()
+    output = embedding.embed_query(document)
+    assert len(output) > 0

libs/cli/langchain_cli/integration_template/tests/unit_tests/test_embeddings.py

@@ -0,0 +1,9 @@
+"""Test embedding model integration."""
+
+
+from __module_name__.embeddings import __ModuleName__Embeddings
+
+
+def test_initialization() -> None:
+    """Test embedding model initialization."""
+    __ModuleName__Embeddings()

libs/cli/langchain_cli/integration_template/tests/unit_tests/test_imports.py

@@ -1,6 +1,11 @@
 from __module_name__ import __all__
 
-EXPECTED_ALL = ["__ModuleName__LLM", "Chat__ModuleName__", "__ModuleName__VectorStore"]
+EXPECTED_ALL = [
+    "__ModuleName__LLM",
+    "Chat__ModuleName__",
+    "__ModuleName__VectorStore",
+    "__ModuleName__Embeddings",
+]
 
 
 def test_all_imports() -> None: