Cleaning up utils, adding error handling

Author: ajosh0504

data/mongodb_docs_embeddings.json

@@ -80,7 +80,7 @@
    "outputs": [],
    "source": [
     "# Obtain a Voyage API key from our AI model proxy and set it as an environment variable-- DO NOT CHANGE\n",
-    "set_env(\"VOYAGE_API_KEY\", PASSKEY)"
+    "set_env([\"voyageai\"], PASSKEY)"
    ]
   },
   {
@@ -764,8 +764,12 @@
     "    messages = <CODE_BLOCK_17>\n",
     "    # Send the chat messages to the AI model proxy get back an LLM response\n",
     "    response = requests.post(url=SERVERLESS_URL, json={\"task\": \"completion\", \"data\": messages})\n",
-    "    # Print the final answer\n",
-    "    print(response.json()[\"text\"])"
+    "    if response.status_code == 200:\n",
+    "        # Print the final answer\n",
+    "        print(response.json()[\"text\"])\n",
+    "    else:\n",
+    "        # Print the error message\n",
+    "        print(response.json()[\"error\"])"
    ]
   },
   {

labs/ai-rag-lab.ipynb

@@ -70,7 +70,7 @@
    "outputs": [],
    "source": [
     "# Obtain a Voyage API key from our AI model proxy and set it as an environment variable-- DO NOT CHANGE\n",
-    "set_env(\"VOYAGE_API_KEY\", PASSKEY)"
+    "set_env([\"voyageai\"], PASSKEY)"
    ]
   },
   {

labs/vector-search-lab.ipynb

@@ -80,7 +80,7 @@
    "outputs": [],
    "source": [
     "# Obtain a Voyage API key from our AI model proxy and set it as an environment variable-- DO NOT CHANGE\n",
-    "set_env(\"VOYAGE_API_KEY\", PASSKEY)"
+    "set_env([\"voyageai\"], PASSKEY)"
    ]
   },
   {
@@ -853,8 +853,12 @@
     "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
     "    # Send the chat messages to the AI model proxy get back an LLM response\n",
     "    response = requests.post(url=SERVERLESS_URL, json={\"task\": \"completion\", \"data\": messages})\n",
-    "    # Print the final answer\n",
-    "    print(response.json()[\"text\"])"
+    "    if response.status_code == 200:\n",
+    "        # Print the final answer\n",
+    "        print(response.json()[\"text\"])\n",
+    "    else:\n",
+    "        # Print the error message\n",
+    "        print(response.json()[\"error\"])"
    ]
   },
   {

solutions/ai-rag-lab.ipynb

@@ -70,7 +70,7 @@
    "outputs": [],
    "source": [
     "# Obtain a Voyage API key from our AI model proxy and set it as an environment variable-- DO NOT CHANGE\n",
-    "set_env(\"VOYAGE_API_KEY\", PASSKEY)"
+    "set_env([\"voyageai\"], PASSKEY)"
    ]
   },
   {

solutions/vector-search-lab.ipynb

@@ -1,12 +1,15 @@
 from pymongo.errors import OperationFailure
 from pymongo.collection import Collection
+from langchain_aws import ChatBedrock
+from langchain_openai import AzureChatOpenAI
+from langchain_google_genai import ChatGoogleGenerativeAI
 import requests
-from typing import Dict
+from typing import Dict, List
 import time
 import os
 
 SLEEP_TIMER = 5
-SERVERLESS_URL = os.getenv("SERVERLESS_URL")
+SERVERLESS_URL = "https://vtqjvgchmwcjwsrela2oyhlegu0hwqnw.lambda-url.us-west-2.on.aws/"
 SANDBOX_NAME = os.getenv("CODESPACE_NAME") or os.getenv("_SANDBOX_ID")
 
 
@@ -85,17 +88,48 @@ def track_progress(task: str, workshop_id: str) -> None:
     payload = {"task": task, "workshop_id": workshop_id, "sandbox_id": SANDBOX_NAME}
     requests.post(url=SERVERLESS_URL, json={"task": "track_progress", "data": payload})
 
-def set_env(name: str, passkey: str) -> None:
+
+def set_env(providers: List[str], passkey: str) -> None:
     """
-    Set environment variable in sandbox
+    Set environment variables in sandbox
 
     Args:
-        name (str): Environment variable name
+        providers (List[str]): List of provider names
         passkey (str): Passkey to get token
     """
-    response = requests.post(url=SERVERLESS_URL, json={"task": "get_token", "data": passkey})
-    if response.status_code == 200:
-        os.environ[name] = response.json().get("token")
-        print(f"{name} environment variable set successfully.")
+    for provider in providers:
+        payload = {"provider": provider, "passkey": passkey}
+        response = requests.post(url=SERVERLESS_URL, json={"task": "get_token", "data": payload})
+        status_code = response.status_code
+        if status_code == 200:
+            result = response.json().get("token")
+            for key in result:
+                os.environ[key] = result[key]
+                print(f"Set {key} environment variable.")
+        elif status_code == 401:
+            print(f"{response.json()['error']}. Follow steps in the lab documentation to obtain your own credentials and set them as environment variables.")
+        else:
+            print(f"{response.json()['error']}")
+
+
+def get_llm(provider: str):
+    if provider == "aws":
+        return ChatBedrock(
+            model_id="anthropic.claude-3-5-sonnet-20240620-v1:0",
+            model_kwargs=dict(temperature=0),
+            region_name="us-west-2",
+        )
+    elif provider == "google":
+        return ChatGoogleGenerativeAI(
+            model="gemini-1.5-pro",
+            temperature=0,
+        )
+    elif provider == "microsoft":
+        return AzureChatOpenAI(
+            azure_endpoint="https://gai-326.openai.azure.com/",
+            azure_deployment="gpt-4o",
+            api_version="2023-06-01-preview",
+            temperature=0,
+        )
     else:
-        print("Passkey expired. Follow steps in the lab documentation to obtain your own credentials and set them as environment variables.")
+        print("Unsupported provider. provider can be one of 'aws', 'google', 'microsoft'.")