Merge branch 'main' into release-0.0.6.5-blog

.github/workflows/lmql-ci.yml

@@ -0,0 +1,68 @@
+name: Run Tests and Publish Wheel
+
+on:
+  # on release publish
+  release:
+    types: [released]
+
+
+jobs:
+    test-without-hf-transformers:
+      runs-on: lmql-ci
+      steps:
+          - uses: actions/checkout@v3
+          - name: Setup Fresh Virtual Environment
+            run: |
+              pip install --upgrade pip
+              python3.10 -m venv env
+              export PATH=$PATH:/home/docker/.local/bin
+              source env/bin/activate
+              echo "VIRTUAL ENV:" $VIRTUAL_ENV
+          - name: Install Dependencies
+            run: source env/bin/activate && pip install -e . && pip install langchain
+          - name: Greet
+            env:
+              OPENAI_API_KEY: ${{ secrets.LMQL_CI_OPENAI_KEY }}
+            run: source env/bin/activate && python -m lmql.cli hello openai
+    test-with-hf-transformers:
+        runs-on: lmql-ci
+        needs: [test-without-hf-transformers]
+        steps:
+            - uses: actions/checkout@v3
+            - name: Setup Fresh Virtual Environment
+              run: |
+                pip install --upgrade pip
+                python3.10 -m venv env
+                export PATH=$PATH:/home/docker/.local/bin
+                source env/bin/activate
+                echo "VIRTUAL ENV:" $VIRTUAL_ENV
+            - name: Install Dependencies
+              run: source env/bin/activate && pip install -e '.[hf]' && pip install langchain
+            - name: Run Tests
+              env:
+                OPENAI_API_KEY: ${{ secrets.LMQL_CI_OPENAI_KEY }}
+              run: source env/bin/activate && python src/lmql/tests/all.py --failearly
+    publish:
+      runs-on: lmql-ci
+      needs: [test-with-hf-transformers, test-without-hf-transformers]
+      steps:
+          - uses: actions/checkout@v3
+          - name: Setup Fresh Virtual Environment
+            run: |
+              pip install --upgrade pip
+              python3.10 -m venv env
+              export PATH=$PATH:/home/docker/.local/bin
+              source env/bin/activate
+              echo "VIRTUAL ENV:" $VIRTUAL_ENV  
+          - name: Install Packaging Dependencies
+            run: pip install build twine
+          - name: Package
+            env:
+              VERSION: ${{ github.ref }}
+            run: bash scripts/wheel.sh $(echo $VERSION | sed 's/^refs\/tags\/v//')
+          - name: Publish
+            env:
+              TWINE_USERNAME: __token__
+              TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
+              VERSION: ${{ github.ref }}
+            run: bash scripts/pypi-release.sh lmql-$(echo $VERSION | sed 's/^refs\/tags\/v//') # --production

.github/workflows/lmql-tests.yml

@@ -0,0 +1,40 @@
+name: Run Tests
+
+on: workflow_dispatch  
+
+jobs:
+    test-without-hf-transformers:
+      runs-on: lmql-ci
+      steps:
+          - uses: actions/checkout@v3
+          - name: Setup Fresh Virtual Environment
+            run: |
+              pip install --upgrade pip
+              python3.10 -m venv env
+              export PATH=$PATH:/home/docker/.local/bin
+              source env/bin/activate
+              echo "VIRTUAL ENV:" $VIRTUAL_ENV
+          - name: Install Dependencies
+            run: source env/bin/activate && pip install -e . && pip install langchain
+          - name: Greet
+            env:
+              OPENAI_API_KEY: ${{ secrets.LMQL_CI_OPENAI_KEY }}
+            run: source env/bin/activate && python -m lmql.cli hello openai
+    test-with-hf-transformers:
+        runs-on: lmql-ci
+        needs: [test-without-hf-transformers]
+        steps:
+            - uses: actions/checkout@v3
+            - name: Setup Fresh Virtual Environment
+              run: |
+                pip install --upgrade pip
+                python3.10 -m venv env
+                export PATH=$PATH:/home/docker/.local/bin
+                source env/bin/activate
+                echo "VIRTUAL ENV:" $VIRTUAL_ENV
+            - name: Install Dependencies
+              run: source env/bin/activate && pip install -e '.[hf]' && pip install langchain
+            - name: Run Tests
+              env:
+                OPENAI_API_KEY: ${{ secrets.LMQL_CI_OPENAI_KEY }}
+              run: source env/bin/activate && python src/lmql/tests/all.py --failearly

docs/source/language/azure.md

@@ -4,32 +4,48 @@ LMQL also supports OpenAI models hosted on Azure. To use these models, you need
 
 ## Configuration via Environment Variables
 
-To configure an LMQL runtime as a whole to use Azure endpoints for OpenAI models, you can provide the following environment variables:
+To configure an LMQL runtime as a whole to use a specific Azure deployed model for OpenAI calls, you can provide the following environment variables:
 
 ```
-OPENAI_API_TYPE=azure
-AZURE_OPENAI_<ENV_NAME>_ENDPOINT=<endpoint>
-AZURE_OPENAI_<ENV_NAME>_KEY=<key>
-```
+# set the API type based on whether you want to use a completion or chat endpoint
+OPENAI_API_TYPE=azure|azure-chat 
 
-where `<ENV_NAME>` is the name of the hosted OpenAI model, e.g. `gpt-3.5-turbo` but uppercase and with `-` replaced by `_`. For example, to configure the `gpt-3.5-turbo` model, you would use the `AZURE_OPENAI_GPT-3_5-TURBO_ENDPOINT` and `AZURE_OPENAI_GPT-3_5-TURBO_KEY` environment variables.
+# your Azure API base URL, e.g. 'https://<YOUR_BASE>.openai.azure.com/'
+OPENAI_API_BASE=<API_BASE> 
 
-## Configuration via `lmql.model`
+# set your API key, can also be provided per deployment 
+# via OPENAI_API_KEY_{<your-deployment-name>.upper()}
+OPENAI_API_KEY=<key>
+```
 
-If you prefer to configure Azure credentials on a per-query basis, you can also specify the endpoint and API key to use as part of a `lmql.model(...)` expression:
+When using your Azure models, make sure to invoke them as `openai/<DEPLOYMENT NAME>` in your query code. If you need more control, or want to use different deployments, base URLs or api versions across your application, please refer to the next section.
 
-```{lmql}
+## Configuration via `lmql.model`
 
-name::azure-api-key
+If you need to configure Azure credentials on a per-query basis, you can also specify the Azure access credentials as part of an `lmql.model(...)` object:
+
+```python
+my_azure_model = lmql.model(
+    # the name of your deployed model/engine, e.g. 'my-model'
+    "openai/<DEPLOYMENT>", 
+    # set to 'azure-chat' for chat endpoints and 'azure' for completion endpoints
+    api_type="azure|azure-chat",  
+    # your Azure API base URL, e.g. 'https://<YOUR_BASE>.openai.azure.com/'
+    api_base="<API_BASE>", 
+    # your API key, can also be provided via env variable OPENAI_API_KEY 
+    # or OPENAI_API_KEY_{<your-deployment-name>.upper()}
+    [api_key="<API_KEY>"] , 
+    # API version, defaults to '2023-05-15'
+    [api_version="API_VERSION",]
+    # prints the full endpoint URL to stdout on each query (alternatively OPENAI_VERBOSE=1)
+    [verbose=False] 
+)
+```
 
-import lmql
+You can then use this model in your query code as follows:
 
-argmax
-    "Hello[WHO]"
-from
-    lmql.model("openai/gpt-3.5-turbo", endpoint="azure://<ENDPOINT>", azure_api_key="<API_KEY>")
-where
-    STOPS_AT(WHO, "\n") and len(TOKENS(WHO)) < 10
-```
+```{lmql}
+name::use-azure-model
 
-Where `<ENDPOINT>` should not include the `https://` prefix and `<API_KEY>` is the API key for the deployed model.
+argmax "Hello [WHO]" from my_azure_model
+```

scripts/pypi-release.sh

@@ -1,13 +1,13 @@
 if [ "$2" == "--production" ]; then
     echo "[WARNING] Uploading to production PyPI"
     echo "Source Distribution"
-    python -m twine upload dist/$1.tar.gz
+    python -m twine upload dist/$1.tar.gz -u $TWINE_USERNAME -p $TWINE_PASSWORD
     echo "Wheel Distribution"
-    python -m twine upload dist/$1-py3-none-any.whl
+    python -m twine upload dist/$1-py3-none-any.whl -u $TWINE_USERNAME -p $TWINE_PASSWORD
 else
     echo "Uploading to test.pypi.org"
     echo "Source Distribution"
-    python -m twine upload --repository testpypi dist/$1.tar.gz
+    python -m twine upload --repository testpypi dist/$1.tar.gz -u $TWINE_USERNAME -p $TWINE_PASSWORD
     echo "Wheel Distribution"
-    python -m twine upload --repository testpypi dist/$1-py3-none-any.whl
+    python -m twine upload --repository testpypi dist/$1-py3-none-any.whl -u $TWINE_USERNAME -p $TWINE_PASSWORD
 fi

scripts/wheel.sh

@@ -1,3 +1,6 @@
+# fail if any commands fails
+set -e
+
 COMMIT=$(git rev-parse HEAD)
 HAS_UNSTAGED=$(git diff-index --quiet HEAD -- src; echo $?)
 

src/lmql/cli.py

@@ -173,7 +173,7 @@ def hello():
 
     if backend is None or backend == "hf":
         code_local = """
-    argmax "Hello[WHO]" from "local:gpt2-medium" where len(WHO) < 10    
+    argmax "Hello[WHO]" from "local:gpt2-medium" where len(TOKENS(WHO)) < 10
     """
         print("[Greeting 🤗 Transformers]")
         asyncio.run(lmql.run(code_local, output_writer=lmql.printing))
@@ -182,10 +182,8 @@ def hello():
         import lmql.runtime.dclib as dc
         dc.clear_tokenizer()
         print("[Greeting OpenAI]")
-        code_openai = """
-    argmax "Hello[WHO]" from "openai/text-ada-001" where len(WHO) < 10    
-    """
-        asyncio.run(lmql.run(code_openai, output_writer=lmql.printing))
+        code_openai = 'argmax "Hello[WHO]" from "openai/text-ada-001" where len(TOKENS(WHO)) < 10 and not "\\n" in WHO'
+        asyncio.run(lmql.run(code_openai, output_writer=lmql.printing, model="openai/text-ada-001"))
 
 def basic_samples():
     from lmql.tests.test_sample_queries import main

src/lmql/runtime/bopenai/openai_api.py

@@ -53,8 +53,15 @@ def __enter__(self):
     def __exit__(self, *args):
         self.task.cancel()
 
+def is_azure_chat(kwargs):
+    if not "api_config" in kwargs: return False
+    api_config = kwargs["api_config"]
+    if not "api_type" in api_config: 
+        return os.environ.get("OPENAI_API_TYPE", "azure") == "azure-chat"
+    return ("api_type" in api_config and "azure-chat" in api_config.get("api_type", ""))
+
 async def complete(**kwargs):
-    if kwargs["model"].startswith("gpt-3.5-turbo") or "gpt-4" in kwargs["model"]:
+    if kwargs["model"].startswith("gpt-3.5-turbo") or "gpt-4" in kwargs["model"] or is_azure_chat(kwargs):
         async for r in chat_api(**kwargs): yield r
     else:
         async for r in completion_api(**kwargs): yield r
@@ -97,36 +104,40 @@ def tagged_segments(s):
 
 def get_azure_config(model, api_config):
     endpoint = api_config.get("endpoint", None)
+    api_type = api_config.get("api_type", os.environ.get("OPENAI_API_TYPE", ""))
 
-    # first check endpoint configuration
-    if endpoint is not None and endpoint.startswith("azure:"):
-        endpoint = "https:" + endpoint[6:]
-        env_name = model.upper().replace(".", "_")
+    if (api_type == "azure" or api_type == "azure-chat"):
+        api_base = os.environ.get("OPENAI_API_BASE", api_config.get("api_base", None))
+        assert api_base is not None, "Please specify the Azure API base URL as 'api_base' or environment variable OPENAI_API_BASE"
+        api_version = os.environ.get("OPENAI_API_VERSION", api_config.get("api_version", "2023-05-15"))
+        deployment = os.environ.get("OPENAI_DEPLOYMENT", api_config.get("api_deployment", model))
 
-        azure_key = api_config.get("azure_api_key", None)
-        if azure_key is None:
-            assert "AZURE_OPENAI_" + env_name + "_KEY" in os.environ, f"Please set the environment variable AZURE_OPENAI_{env_name}_KEY to your Azure API key, or specify it in code as 'lmql.model(..., azure_api_key=...')"
-            azure_key = os.environ["AZURE_OPENAI_" + env_name + "_KEY"]
+        deployment_specific_api_key = f"OPENAI_API_KEY_{deployment.upper()}"
+        api_key = os.environ.get(deployment_specific_api_key, None) or os.environ.get("OPENAI_API_KEY", api_config.get("api_key", None))
+        assert api_key is not None, "Please specify the Azure API key as 'api_key' or environment variable OPENAI_API_KEY or OPENAI_API_KEY_<DEPLOYMENT>"
 
-        headers = {
-            "Content-Type": "application/json",
-            "api-key": azure_key,
-        }
+        is_chat = api_type == "azure-chat"
 
-        return endpoint, headers
+        if is_chat:
+            endpoint = f"{api_base}/openai/deployments/{deployment}/chat/completions"
+        else:
+            endpoint = f"{api_base}/openai/deployments/{deployment}/completions"
+
+        if api_version is not None:
+            endpoint += f"?api-version={api_version}"
 
-    # then check env
-    if os.environ.get("OPENAI_API_TYPE", 'openai') == 'azure':
-        model_env_name = model.upper().replace(".", "_")
-        endpoint = os.environ[f"AZURE_OPENAI_{model_env_name}_ENDPOINT"]
-        assert "AZURE_OPENAI_" + model_env_name + "_KEY" in os.environ, f"Please set the environment variable AZURE_OPENAI_{model_env_name}_KEY to your Azure API key, or specify it in code as 'lmql.model(..., azure_api_key=...')"
-        key = os.environ[f"AZURE_OPENAI_{model_env_name}_KEY"]
         headers = {
             "Content-Type": "application/json",
-            "api-key": key,
+            "api-key": api_key,
         }
+
+        if "verbose" in api_config and api_config["verbose"] or os.environ.get("OPENAI_VERBOSE", "0") == "1":
+            print(f"Using Azure API endpoint: {endpoint}", is_chat)
+
         return endpoint, headers
 
+    return None
+
 def get_endpoint_and_headers(kwargs):
     model = kwargs["model"]
     api_config = kwargs.pop("api_config", {})
@@ -139,9 +150,9 @@ def get_endpoint_and_headers(kwargs):
 
     # otherwise use custom endpoint as plain URL without authorization
     if endpoint is not None:
-        if not endpoint.endswith("http"):
+        if not endpoint.startswith("http"):
             endpoint = f"http://{endpoint}"
-        return endpoint + "/completions", {
+        return endpoint, {
             "Content-Type": "application/json"
         }
 
@@ -221,7 +232,7 @@ async def chat_api(**kwargs):
 
         current_chunk = ""
         stream_start = time.time()
-        
+
         async with aiohttp.ClientSession() as session:
             endpoint, headers = get_endpoint_and_headers(kwargs)
             async with session.post(

src/lmql/tests/all.py

@@ -1,21 +1,31 @@
 import os
 import sys
+import subprocess
 
 THIS_DIR = os.path.dirname(__file__)
 files = sorted(os.listdir(THIS_DIR))
+TEST_TIMEOUT = float(os.environ.get("TEST_TIMEOUT", 3*60.0))
 
 errors = 0 
 files = [f for f in files if f.startswith("test_")]
 
 for i,f in enumerate(files):
     try:
         print(">", f"[{i+1}/{len(files)}]", f)
-        result = os.system("python " + os.path.join(THIS_DIR, f))
+
+        cmd = "python " + os.path.join(THIS_DIR, f)
+        timeout = TEST_TIMEOUT
+        result = subprocess.call(cmd, shell=True, timeout=timeout)
 
         if result == 2:
             raise KeyboardInterrupt
         if result != 0:
             errors += 1
+            if "--failearly" in sys.argv:
+                break
+    except subprocess.TimeoutExpired:
+        print(">", f"[{i+1}/{len(files)}]", f, "timed out after", timeout, "seconds")
+        sys.exit(1)
 
     except KeyboardInterrupt:
         sys.exit(1)

src/lmql/tests/expr_test_utils.py

@@ -160,9 +160,9 @@ def run_all_tests(g):
                     g[k]()
                 termcolor.cprint("OK", "green")
         except AssertionError as e:
-            print(e)
             num_errors += 1
             termcolor.cprint("FAILED", "red")
+            print(e)
 
     # wait for all tasks to finish
     try:
@@ -181,7 +181,6 @@ def run_all_tests(g):
         print(num_errors, "test(s) failed.")
         sys.exit(1)
     else: 
-        print("All tests passed.")
         sys.exit(0)
 
 def enable_show_transformed():