Feature/googleaistudio (#121)

* Added jsonnet for GoogleAIStudio * Port the template to 0.13, env secrets for k8s --------- Co-authored-by: JackColquitt <[email protected]>
trustgraph-ai · Oct 15, 2024 · af0c854 · af0c854
1 parent ec444d1
commit af0c854
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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",
+}