Merge pull request #36 from arpitjain099/alert-autofix-18

Fix code scanning alert no. 18: Full server-side request forgery
arpitjain099 · Oct 20, 2024 · 1dac837 · 1dac837
2 parents 98b4f8e + ddc261a
commit 1dac837
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/End_to_end_Solutions/InsightsGenerator/insights_generator/core/OAI_client.py b/End_to_end_Solutions/InsightsGenerator/insights_generator/core/OAI_client.py
@@ -5,8 +5,11 @@
 import tiktoken
 
 def make_prompt_request(prompt, max_tokens = 2048, timeout = 4):
-    #url = "https://api.openai.com/v1/embeddings"
+    # Whitelist of allowed URLs
+    allowed_urls = ["https://api.openai.com/v1/embeddings", "https://another-trusted-url.com"]
     url = os.getenv("AOAI_ENDPOINT")
+    if url not in allowed_urls:
+        raise ValueError("The provided URL is not allowed.")
     key = os.getenv("AOAI_KEY")
 
     payload_dict = {"prompt": prompt,