From 37ced42f5639c89226a82199c2cbf18b39b9525f Mon Sep 17 00:00:00 2001
From: Cyrus Pellet <cyrus135@hotmail.fr>
Date: Tue, 9 Jan 2024 13:16:07 +0100
Subject: [PATCH] Added OllamaEmbeddings component with documentation

---
 docs/docs/components/embeddings.mdx           | 20 ++++++---
 .../components/embeddings/OllamaEmbeddings.py | 41 +++++++++++++++++++
 src/backend/langflow/config.yaml              |  2 +
 3 files changed, 58 insertions(+), 5 deletions(-)
 create mode 100644 src/backend/langflow/components/embeddings/OllamaEmbeddings.py

diff --git a/docs/docs/components/embeddings.mdx b/docs/docs/components/embeddings.mdx
index 1b7ca1dbe..015ba1ce1 100644
--- a/docs/docs/components/embeddings.mdx
+++ b/docs/docs/components/embeddings.mdx
@@ -1,11 +1,13 @@
-import Admonition from '@theme/Admonition';
+import Admonition from "@theme/Admonition";
 
 # Embeddings
 
 <Admonition type="caution" icon="🚧" title="ZONE UNDER CONSTRUCTION">
-    <p>
-        We appreciate your understanding as we polish our documentation – it may contain some rough edges. Share your feedback or report issues to help us improve! 🛠️📝
-    </p>
+  <p>
+    We appreciate your understanding as we polish our documentation – it may
+    contain some rough edges. Share your feedback or report issues to help us
+    improve! 🛠️📝
+  </p>
 </Admonition>
 
 Embeddings are vector representations of text that capture the semantic meaning of the text. They are created using text embedding models and allow us to think about the text in a vector space, enabling us to perform tasks like semantic search, where we look for pieces of text that are most similar in the vector space.
@@ -110,4 +112,12 @@ Vertex AI is a cloud computing platform offered by Google Cloud Platform (GCP).
 - **top_k:** How the model selects tokens for output, the next token is selected from – defaults to `40`.
 - **top_p:** Tokens are selected from most probable to least until the sum of their – defaults to `0.95`.
 - **tuned_model_name:** The name of a tuned model. If provided, model_name is ignored.
-- **verbose:** This parameter is used to control the level of detail in the output of the chain. When set to True, it will print out some internal states of the chain while it is being run, which can help debug and understand the chain's behavior. If set to False, it will suppress the verbose output – defaults to `False`.
\ No newline at end of file
+- **verbose:** This parameter is used to control the level of detail in the output of the chain. When set to True, it will print out some internal states of the chain while it is being run, which can help debug and understand the chain's behavior. If set to False, it will suppress the verbose output – defaults to `False`.
+
+### OllamaEmbeddings
+
+Used to load [Ollama’s](https://ollama.ai/) embedding models. Wrapper around LangChain's [Ollama API](https://python.langchain.com/docs/integrations/text_embedding/ollama).
+
+- **model** The name of the Ollama model to use – defaults to `llama2`.
+- **base_url** The base URL for the Ollama API – defaults to `http://localhost:11434`.
+- **temperature** Tunes the degree of randomness in text generations. Should be a non-negative value – defaults to `0`.
diff --git a/src/backend/langflow/components/embeddings/OllamaEmbeddings.py b/src/backend/langflow/components/embeddings/OllamaEmbeddings.py
new file mode 100644
index 000000000..388afe3f6
--- /dev/null
+++ b/src/backend/langflow/components/embeddings/OllamaEmbeddings.py
@@ -0,0 +1,41 @@
+from typing import Optional
+
+from langflow import CustomComponent
+from langchain.embeddings.base import Embeddings
+from langchain_community.embeddings import OllamaEmbeddings
+
+class OllamaEmbeddingsComponent(CustomComponent):
+    """
+    A custom component for implementing an Embeddings Model using Ollama.
+    """
+
+    display_name: str = "Ollama Embeddings"
+    description: str = "Embeddings model from Ollama."
+    documentation = "https://python.langchain.com/docs/integrations/text_embedding/ollama"
+    beta = True
+
+    def build_config(self):
+        return {
+            "model": {
+                "display_name": "Ollama Model",
+            },
+            "base_url": {"display_name": "Ollama Base URL"},
+            "temperature": {"display_name": "Model Temperature"},
+            "code": {"show": False},
+        }
+    
+    def build(
+        self,
+        model: str = "llama2",
+        base_url: str = "http://localhost:11434",
+        temperature: Optional[float] = None,
+    ) -> Embeddings:
+        try:
+            output = OllamaEmbeddings(
+                model=model,
+                base_url=base_url,
+                temperature=temperature
+            )  # type: ignore
+        except Exception as e:
+            raise ValueError("Could not connect to Ollama API.") from e
+        return output
\ No newline at end of file
diff --git a/src/backend/langflow/config.yaml b/src/backend/langflow/config.yaml
index 816f4738c..caabffbf9 100644
--- a/src/backend/langflow/config.yaml
+++ b/src/backend/langflow/config.yaml
@@ -106,6 +106,8 @@ embeddings:
     documentation: "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/google_vertex_ai_palm"
   AmazonBedrockEmbeddings:
     documentation: "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/bedrock"
+  OllamaEmbeddings:
+    documentation: "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/ollama"
 
 llms:
   OpenAI: