Implement openai.embeddings.create

.github/workflows/pypi.yml

@@ -0,0 +1,36 @@
+name: Publish to PyPi
+
+on:
+  release:
+    types: [published]
+  workflow_dispatch:
+
+jobs:
+  pypi:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+
+    - uses: actions/setup-python@v2
+      with:
+        python-version: "3.11"
+
+    - name: Install dependencies
+      run: |
+        python3 -m pip install --upgrade pip
+        python3 -m pip install build twine
+        python3 -m pip install -r requirements.txt
+
+    - name: Build Function LLM
+      run: python3 -m build
+
+    - name: Publish to PyPi
+      run: python3 -m twine upload dist/*
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
+
+    - uses: actions/upload-artifact@v4
+      with:
+          name: Wheels
+          path: dist/

.github/workflows/testpypi.yml

@@ -0,0 +1,34 @@
+name: Publish to TestPyPi
+
+on: workflow_dispatch
+
+jobs:
+  pypi:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+
+    - uses: actions/setup-python@v2
+      with:
+        python-version: "3.11"
+
+    - name: Install dependencies
+      run: |
+        python3 -m pip install --upgrade pip
+        python3 -m pip install build twine
+        python3 -m pip install -r requirements.txt
+
+    - name: Build Function LLM
+      run: python3 -m build
+
+    - name: Publish to TestPyPi
+      run: python3 -m twine upload dist/*
+      env:
+        TWINE_REPOSITORY: testpypi
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.TEST_PYPI_TOKEN }}
+
+    - uses: actions/upload-artifact@v4
+      with:
+          name: Wheels
+          path: dist/

fxn_llm/__init__.py

@@ -3,4 +3,5 @@
 #   Copyright © 2024 NatML Inc. All Rights Reserved.
 #
 
+from .locally import locally, LLMClient, LLMProvider
 from .version import *

fxn_llm/locally.py

@@ -0,0 +1,68 @@
+# 
+#   Function
+#   Copyright © 2024 NatML Inc. All Rights Reserved.
+#
+
+from fxn import Function
+from numpy import float32
+from numpy.typing import NDArray
+from types import MethodType
+from typing import List, Literal, Optional, TypeVar
+
+LLMProvider = Literal["openai", "anthropic"]
+LLMClient = TypeVar("LLMClient")
+
+def locally (
+    client: LLMClient,
+    provider: LLMProvider="openai",
+    access_key: str=None,
+    api_url: str=None,
+    fxn: Function=None
+) -> LLMClient:
+    """
+    Patch your LLM client to run locally.
+
+    Parameters:
+        client (OpenAI | Anthropic): LLM provider client.
+        provider (LLMProvider): LLM provider identifier. Defaults to `openai`.
+        access_key (str): Function access key.
+        api_url (str): Function API URL.
+        fxn (Function): Function client.
+
+    Returns:
+        OpenAI | Anthropic: LLM provider client patched to run locally.
+    """
+    fxn = fxn if fxn is not None else Function(access_key=access_key, api_url=api_url)
+    if provider == "openai":
+        from openai.types import CreateEmbeddingResponse, Embedding
+        from openai.types.create_embedding_response import Usage
+        def embeddings_create (
+            self,
+            *,
+            input: str | List[str],
+            model: str,
+            dimensions: Optional[int]=None,
+            encoding_format: Optional[Literal["float", "base64"]]=None,
+            **kwargs
+        ) -> CreateEmbeddingResponse:
+            encoding_format = encoding_format if encoding_format is not None else "float"
+            assert dimensions is None, "Explicit dimensionality is not yet supported"
+            assert encoding_format == "float", "Base64 encoding format is not yet supported"
+            input = [input] if isinstance(input, str) else input
+            prediction = fxn.predictions.create(tag=model, inputs={ "input": input })
+            embeddings: NDArray[float32] = prediction.results[0]
+            return CreateEmbeddingResponse(
+                data=[Embedding(
+                    embedding=data.tolist(),
+                    index=idx,
+                    object="embedding"
+                ) for idx, data in enumerate(embeddings)],
+                model=model,
+                object="list",
+                usage=Usage(prompt_tokens=0, total_tokens=0)
+            )
+        client.embeddings.create = MethodType(embeddings_create, client.embeddings)
+        return client
+    elif provider == "anthropic":
+        pass
+    return client

requirements.txt

@@ -1,3 +1,5 @@
 anthropic
 fxn
-openai
+openai
+pytest
+python_dotenv

test/openai_test.py

@@ -3,3 +3,15 @@
 #   Copyright © 2024 NatML Inc. All Rights Reserved.
 #
 
+from fxn_llm import locally
+from openai import OpenAI
+from openai.types import CreateEmbeddingResponse
+
+def test_create_embeddings ():
+    openai = OpenAI(api_key="fxn")
+    openai = locally(openai)
+    embedding = openai.embeddings.create(
+        model="@nomic/nomic-embed-text-v1.5-quant",
+        input="search_query: What is the capital of France?"
+    )
+    assert isinstance(embedding, CreateEmbeddingResponse), f"Returned embeddings has invalid type: {type(embedding)}"