📦 NEW: Embed

Author: arre-ankit

examples/embed/embed.py

@@ -0,0 +1,25 @@
+# Experimental upcoming beta AI primitve.
+# Please refer to the documentation for more information: https://langbase.com/docs for more information.
+import os
+from dotenv import load_dotenv
+from langbase import Langbase
+
+load_dotenv()
+
+# Cconfigure the Langbase client with your API key
+langbase = Langbase(api_key=os.environ.get("LANGBASE_API_KEY"))
+
+def main():
+    """
+    Generates embeddings for the given text chunks.
+    """
+    response = langbase.embed(
+        chunks=[
+            "Langbase is the most powerful serverless platform for building AI agents with memory. Build, scale, and evaluate AI agents with semantic memory (RAG) and world-class developer experience. We process billions of AI messages/tokens daily. Built for every developer, not just AI/ML experts."
+        ],
+        embedding_model="openai:text-embedding-3-large",
+    )
+    print(response)
+
+if __name__ == "__main__":
+    main()

langbase/types.py

@@ -339,6 +339,14 @@ class ToolWebSearchOptions(TypedDict, total=False):
     domains: List[str]
     api_key: str
 
+class EmbedOptions(TypedDict, total=False):
+    """Options for embedding generation."""
+    chunks: List[str]
+    embedding_model: EmbeddingModel
+
+
+EmbedResponse = List[List[float]]
+
 
 class ToolWebSearchResponse(TypedDict):
     """Response from web search."""

tests/test_langbase.py

@@ -249,17 +249,33 @@ def test_threads_messages_list(self, mock_get):
         mock_get.assert_called_once_with("/v1/threads/thread_123/messages")
         self.assertEqual(result, [{"id": "msg_123", "content": "Hello"}])
 
-
-
     @patch("langbase.request.Request.post")
     def test_embed(self, mock_post):
         """Test embed method."""
         mock_post.return_value = [[0.1, 0.2, 0.3]]
-        result = self.lb.embed(
+        
+        # Test with embedding model
+        result_with_model = self.lb.embed(
+            chunks=["Test text"],
+            embedding_model="test-model"
+        )
+        
+        mock_post.assert_called_with(
+            "/v1/embed",
+            {"chunks": ["Test text"], "embeddingModel": "test-model"}
+        )
+        self.assertEqual(result_with_model, [[0.1, 0.2, 0.3]])
+
+        # Test without embedding model
+        result_without_model = self.lb.embed(
             chunks=["Test text"]
         )
-        mock_post.assert_called_once()
-        self.assertEqual(result, [[0.1, 0.2, 0.3]])
+        
+        mock_post.assert_called_with(
+            "/v1/embed",
+            {"chunks": ["Test text"]}
+        )
+        self.assertEqual(result_without_model, [[0.1, 0.2, 0.3]])
 
     @patch("langbase.request.Request.post")
     def test_chunker(self, mock_post):