From af0c8540e35fb45ff97a557f6a01b0ac3aa981f8 Mon Sep 17 00:00:00 2001
From: cybermaggedon <cybermaggedon@gmail.com>
Date: Tue, 15 Oct 2024 20:41:32 +0100
Subject: [PATCH] Feature/googleaistudio (#121)

* Added jsonnet for GoogleAIStudio
* Port the template to 0.13, env secrets for k8s

---------

Co-authored-by: JackColquitt <daniel@kalntera.ai>
---
 templates/components.jsonnet                  |  1 +
 templates/components/googleaistudio.jsonnet   | 81 +++++++++++++++++++
 templates/generate-all                        |  2 +-
 templates/patterns/llm-googleaistudio.jsonnet | 32 ++++++++
 4 files changed, 115 insertions(+), 1 deletion(-)
 create mode 100644 templates/components/googleaistudio.jsonnet
 create mode 100644 templates/patterns/llm-googleaistudio.jsonnet

diff --git a/templates/components.jsonnet b/templates/components.jsonnet
index 7508740c..f27be5cd 100644
--- a/templates/components.jsonnet
+++ b/templates/components.jsonnet
@@ -7,6 +7,7 @@
    "document-rag": import "components/document-rag.jsonnet",
    "embeddings-hf": import "components/embeddings-hf.jsonnet",
    "embeddings-ollama": import "components/embeddings-ollama.jsonnet",
+   "googleaistudio": import "components/googleaistudio.jsonnet",
    "grafana": import "components/grafana.jsonnet",
    "graph-rag": import "components/graph-rag.jsonnet",
    "triple-store-cassandra": import "components/cassandra.jsonnet",
diff --git a/templates/components/googleaistudio.jsonnet b/templates/components/googleaistudio.jsonnet
new file mode 100644
index 00000000..04249c16
--- /dev/null
+++ b/templates/components/googleaistudio.jsonnet
@@ -0,0 +1,81 @@
+local base = import "base/base.jsonnet";
+local images = import "values/images.jsonnet";
+local url = import "values/url.jsonnet";
+local prompts = import "prompts/mixtral.jsonnet";
+
+{
+
+    "googleaistudio-max-output-tokens":: 4096,
+    "googleaistudio-temperature":: 0.0,
+
+    "text-completion" +: {
+    
+        create:: function(engine)
+
+            local envSecrets = engine.envSecrets("bedrock-credentials")
+                .with_env_var("GOOGLEAISTUDIO_KEY", "googleaistudio-key");
+
+            local container =
+                engine.container("text-completion")
+                    .with_image(images.trustgraph)
+                    .with_command([
+                        "text-completion-googleaistudio",
+                        "-p",
+                        url.pulsar,
+                        "-x",
+                        std.toString($["googleaistudio-max-output-tokens"]),
+                        "-t",
+                        std.toString($["googleaistudio-temperature"]),
+                    ])
+                    .with_env_var_secrets(envSecrets)
+                    .with_limits("0.5", "128M")
+                    .with_reservations("0.1", "128M");
+
+            local containerRag =
+                engine.container("text-completion-rag")
+                    .with_image(images.trustgraph)
+                    .with_command([
+                        "text-completion-googleaistudio",
+                        "-p",
+                        url.pulsar,
+                        "-x",
+                        std.toString($["googleaistudio-max-output-tokens"]),
+                        "-t",
+                        std.toString($["googleaistudio-temperature"]),
+                        "-i",
+                        "non-persistent://tg/request/text-completion-rag",
+                        "-o",
+                        "non-persistent://tg/response/text-completion-rag-response",
+                    ])
+                    .with_env_var_secrets(envSecrets)
+                    .with_limits("0.5", "128M")
+                    .with_reservations("0.1", "128M");
+
+            local containerSet = engine.containers(
+                "text-completion", [ container ]
+            );
+
+            local containerSetRag = engine.containers(
+                "text-completion-rag", [ containerRag ]
+            );
+
+            local service =
+                engine.internalService(containerSet)
+                .with_port(8000, 8000, "metrics");
+
+            local serviceRag =
+                engine.internalService(containerSetRag)
+                .with_port(8000, 8000, "metrics");
+
+            engine.resources([
+                envSecrets,
+                containerSet,
+                containerSetRag,
+                service,
+                serviceRag,
+            ])
+
+    },
+
+} + prompts
+
diff --git a/templates/generate-all b/templates/generate-all
index 8b59ce03..0b403620 100755
--- a/templates/generate-all
+++ b/templates/generate-all
@@ -123,7 +123,7 @@ def generate_all(output, version):
     ]:
         for model in [
                 "azure", "azure-openai", "bedrock", "claude", "cohere",
-                "llamafile", "ollama", "openai", "vertexai"
+                "googleaistudio", "llamafile", "ollama", "openai", "vertexai",
         ]:
             for graph in [ "cassandra", "neo4j" ]:
 
diff --git a/templates/patterns/llm-googleaistudio.jsonnet b/templates/patterns/llm-googleaistudio.jsonnet
new file mode 100644
index 00000000..aa56d347
--- /dev/null
+++ b/templates/patterns/llm-googleaistudio.jsonnet
@@ -0,0 +1,32 @@
+{
+    pattern: {
+	name: "googleaistudio",
+        icon: "🤖💬",
+        title: "Add GoogleAIStudio for text completion",
+	description: "This pattern integrates a GoogleAIStudio LLM service for text completion operations. You need a GoogleAISTudio API key to be able to use this service.",
+        requires: ["pulsar", "trustgraph"],
+        features: ["llm"],
+	args: [
+	    {
+		name: "googleaistudio-max-output-tokens",
+                label: "Maximum output tokens",
+		type: "integer",
+		description: "Limit on number tokens to generate",
+                default: 4096,
+		required: true,
+            },
+	    {
+		name: "googleaistudio-temperature",
+                label: "Temperature",
+		type: "slider",
+		description: "Controlling predictability / creativity balance",
+                min: 0,
+                max: 1,
+                step: 0.05,
+                default: 0.5,
+            },
+	],
+        category: [ "llm" ],
+    },
+    module: "components/googleaistudio.jsonnet",
+}