Merge pull request #31 from DataChefHQ/feature/aws-llm-controller

feat: adding aws controller for llm application

.test.sh

@@ -2,5 +2,5 @@
 set -euo pipefail
 
 python -m mypy ./src
-python -m pyright .
+python -m mypy ./tests
 python -m pytest -vv

devenv.lock

@@ -3,11 +3,11 @@
     "devenv": {
       "locked": {
         "dir": "src/modules",
-        "lastModified": 1722262342,
+        "lastModified": 1727963652,
         "owner": "cachix",
         "repo": "devenv",
-        "rev": "11a1ca0ad80bc172d2efda34ae542494442dcf48",
-        "treeHash": "c1be883f8fad6adb0369cef0ac6e6c9bd7f3ec66",
+        "rev": "cb0052e25dbcc8267b3026160dc73cddaac7d5fd",
+        "treeHash": "4d81a7de8e23f71c47c22a2d33fe63787fc2d2e1",
         "type": "github"
       },
       "original": {
@@ -38,11 +38,11 @@
         "systems": "systems"
       },
       "locked": {
-        "lastModified": 1710146030,
+        "lastModified": 1726560853,
         "owner": "numtide",
         "repo": "flake-utils",
-        "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a",
-        "treeHash": "bd263f021e345cb4a39d80c126ab650bebc3c10c",
+        "rev": "c1dfcf08411b08f6b8615f7d8971a2bfa81d5e8a",
+        "treeHash": "0cbff24102dddab81e4c2940d7b356205c955fc6",
         "type": "github"
       },
       "original": {
@@ -95,11 +95,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1720642556,
+        "lastModified": 1724996935,
         "owner": "nlewo",
         "repo": "nix2container",
-        "rev": "3853e5caf9ad24103b13aa6e0e8bcebb47649fe4",
-        "treeHash": "a9c2f1d3f52f288515ca0fb11f9aed970fd869b6",
+        "rev": "fa6bb0a1159f55d071ba99331355955ae30b3401",
+        "treeHash": "a934d246fadcf8b36d28f3577fad413f5ab3f7d3",
         "type": "github"
       },
       "original": {
@@ -126,11 +126,11 @@
     },
     "nixpkgs-stable": {
       "locked": {
-        "lastModified": 1722221733,
+        "lastModified": 1727907660,
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "12bf09802d77264e441f48e25459c10c93eada2e",
-        "treeHash": "e959ebf2e25b21ec31266bef769b447e4b907916",
+        "rev": "5966581aa04be7eff830b9e1457d56dc70a0b798",
+        "treeHash": "55b5fb46cd5d19fe4690148056c5f013a899d746",
         "type": "github"
       },
       "original": {
@@ -150,11 +150,11 @@
         "nixpkgs-stable": "nixpkgs-stable"
       },
       "locked": {
-        "lastModified": 1721042469,
+        "lastModified": 1727854478,
         "owner": "cachix",
         "repo": "pre-commit-hooks.nix",
-        "rev": "f451c19376071a90d8c58ab1a953c6e9840527fd",
-        "treeHash": "91f40b7a3b9f6886bd77482cba5b5cd890415a2e",
+        "rev": "5f58871c9657b5fc0a7f65670fe2ba99c26c1d79",
+        "treeHash": "9bd8fa1bb0d757e1eb29a1eb0d8da485b57b1b31",
         "type": "github"
       },
       "original": {

devenv.nix

@@ -1,17 +1,9 @@
-{
-  pkgs,
-  lib,
-  config,
-  inputs,
-  ...
-}:
+{ pkgs, lib, config, inputs, ... }:
 
 {
   name = "dmv";
   # https://devenv.sh/basics/
-  env = {
-    GREET = "🛠️ Let's hack ";
-  };
+  env = { GREET = "🛠️ Let's hack "; };
 
   # https://devenv.sh/scripts/
   scripts.hello.exec = "echo $GREET";
@@ -34,13 +26,7 @@
   };
 
   # https://devenv.sh/packages/
-  packages = with pkgs; [
-    nixfmt-rfc-style
-    bat
-    jq
-    tealdeer
-    git
-  ];
+  packages = with pkgs; [ nixfmt-rfc-style bat jq tealdeer git ];
 
   languages = {
     # pyright requires npm
@@ -84,6 +70,13 @@
     yamllint = {
       enable = true;
       settings.preset = "relaxed";
+      settings.configuration = ''
+        ---
+        extends: relaxed
+
+        rules:
+          line-length: disable
+      '';
     };
 
     ruff.enable = true;

examples/advance-aws-llm-with-dspy/Pulumi.yaml

@@ -0,0 +1,7 @@
+name: dspy_application
+runtime:
+  name: python
+  options:
+    toolchain: pip
+    virtualenv: venv
+description: A minimal Dspy application

examples/advance-aws-llm-with-dspy/__main__.py

@@ -0,0 +1,31 @@
+import os
+from controller import AwsDspyController
+from damavand.environment import Environment
+
+controller = AwsDspyController(
+    name="my-dspy",
+    region="eu-west-1",
+)
+
+
+def lambda_handler(event, context):
+    return controller.build_or_run(
+        app_id=event.get("app_id", "default"),
+        question=event.get("question", "What llm are you?"),
+    )
+
+
+if __name__ == "__main__":
+    if controller.environment == Environment.LOCAL:
+        # run the lambda handler locally
+        event = {
+            "app_id": os.environ.get("APP_ID", "default"),
+            "question": os.environ.get("QUESTION", "What llm are you?"),
+        }
+        context = {}
+
+        if response := lambda_handler(event, context):
+            print(response)
+    else:
+        # aws automatically calls lambda_handler
+        pass

examples/advance-aws-llm-with-dspy/applications.py

@@ -0,0 +1,14 @@
+import json
+import dspy
+
+
+def default_question(connection: dspy.OpenAI, question: str) -> dict:
+    dspy.settings.configure(lm=connection)
+    predict = dspy.Predict("question -> answer")
+    answer = predict(question=question)
+
+    return {
+        "statusCode": 200,
+        "headers": {"Content-Type": "application/json"},
+        "body": json.dumps({"response": answer}),
+    }

examples/advance-aws-llm-with-dspy/controller.py

@@ -0,0 +1,64 @@
+import dspy
+from typing import Callable, Optional
+
+from damavand.base.controllers.base_controller import runtime
+from damavand.cloud.aws.controllers.llm import AwsLlmController
+
+import applications
+
+
+API_KEY = "EMPTY"
+
+
+class AwsDspyController(AwsLlmController):
+    def __init__(
+        self,
+        name,
+        region: str,
+        model: Optional[str] = None,
+        tags: dict[str, str] = {},
+        **kwargs,
+    ) -> None:
+        super().__init__(name, region, model, tags, **kwargs)
+        self.applications: dict[str, Callable] = {
+            "default": applications.default_question,
+        }
+
+    @property
+    @runtime
+    def connection(self) -> dspy.OpenAI:
+        """Return the dspy OpenAI model."""
+
+        return dspy.OpenAI(
+            model=self.model_id,
+            api_base=f"{self.base_url}/",
+            api_key=API_KEY,
+            model_type="chat",
+        )
+
+    @runtime
+    def run_application(self, app_id: str, question: str, **kwargs) -> dict:
+        """Run the specified application."""
+
+        return self.applications[app_id](question, **kwargs)
+
+    @runtime
+    def build_or_run(self, **kwargs) -> None | dict:
+        """
+        Build or run the application based on the execution mode.
+
+        Parameters
+        ----------
+        kwargs
+            arguments to be passed to the application. Check the `run_application` method for more information.
+
+        Returns
+        -------
+        None | dict
+            If the execution mode is runtime, return the output of the application otherwise None.
+        """
+
+        if self.is_runtime_execution:
+            self.run_application(**kwargs)
+        else:
+            self.provision()

examples/advance-aws-llm-with-dspy/requirements.txt

@@ -0,0 +1,5 @@
+-e ../../../damavand
+pulumi
+boto3
+dspy-ai
+sagemaker

examples/aws-llm-with-openai-client/Pulumi.yaml

@@ -0,0 +1,7 @@
+name: llm-openai-client
+runtime:
+  name: python
+  options:
+    toolchain: pip
+    virtualenv: venv
+description: A simple llm application that uses OpenAI's client.

examples/aws-llm-with-openai-client/__main__.py

@@ -0,0 +1,28 @@
+from damavand.cloud.aws.controllers.llm import AwsLlmController
+
+
+controller = AwsLlmController(
+    name="my-dspy",
+    region="eu-west-1",
+)
+
+
+def lambda_handler(event, context):
+    question = event.get("question")
+    role = event.get("role", "user")
+
+    response = controller.client.chat.completions.create(
+        model=controller.model_id,
+        messages=[{"role": role, "content": question}],
+    )
+
+    return {
+        "statusCode": 200,
+        "headers": {"Content-Type": "application/json"},
+        "body": response,
+    }
+
+
+if __name__ == "__main__":
+    if not controller.is_runtime_execution:
+        controller.provision()

examples/aws-llm-with-openai-client/requirements.txt

@@ -0,0 +1,5 @@
+-e ../../../damavand
+pulumi
+boto3
+sagemaker
+openai

examples/aws-llm-with-simple-requests/Pulumi.yaml

@@ -0,0 +1,7 @@
+name: simple-llm-application
+runtime:
+  name: python
+  options:
+    toolchain: pip
+    virtualenv: venv
+description: A minimal prompt engineering application that uses open source llm models

examples/aws-llm-with-simple-requests/__main__.py

@@ -0,0 +1,53 @@
+import requests
+import os
+
+from damavand.cloud.aws.controllers.llm import AwsLlmController
+
+
+controller = AwsLlmController(
+    name="my-dspy",
+    region="eu-west-1",
+)
+
+
+def ask(question: str, role: str) -> None:
+    headers = {
+        "Content-Type": "application/json",
+    }
+
+    json_data = {
+        "messages": [
+            {
+                "role": role,
+                "content": question,
+            },
+        ],
+        "parameters": {
+            "max_new_tokens": 400,
+        },
+        "stream": False,
+    }
+
+    return requests.post(
+        controller.chat_completions_url,
+        headers=headers,
+        json=json_data,
+    ).json()
+
+
+def lambda_handler(event, context):
+    question = event.get("question")
+    role = event.get("role", "user")
+    return ask(question, role)
+
+
+if __name__ == "__main__":
+    if not controller.is_runtime_execution:
+        controller.provision()
+    else:
+        event = {
+            "question": os.environ.get("QUESTION", "What is the capital of France?"),
+            "role": os.environ.get("ROLE", "user"),
+        }
+        context = {}
+        print(lambda_handler(event, context))

examples/aws-llm-with-simple-requests/requirements.txt

@@ -0,0 +1,4 @@
+-e ../../../damavand
+pulumi
+boto3
+sagemaker

src/damavand/base/controllers/__init__.py

@@ -1,11 +1,7 @@
 from .base_controller import ApplicationController, runtime, buildtime
-from .object_storage import ObjectStorageController
-from .spark import SparkController
 
 __all__ = [
-    "ApplicationController",
-    "ObjectStorageController",
-    "SparkController",
     "runtime",
     "buildtime",
+    "ApplicationController",
 ]