From 6429350c32a11304749213de6951702ec771c111 Mon Sep 17 00:00:00 2001
From: Vivek Nair <vivek@pesto.app>
Date: Mon, 24 Apr 2023 20:56:31 -0400
Subject: [PATCH] fix: check hostname validity, throw exception if not correct
 (#65)

---
 package/gentrace/providers/getters.py  |  8 ++++++++
 package/tests/test_openai_embedding.py |  4 +++-
 package/tests/test_usage.py            | 16 +++++++++++++++-
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/package/gentrace/providers/getters.py b/package/gentrace/providers/getters.py
index 5d9d95b..781ff98 100644
--- a/package/gentrace/providers/getters.py
+++ b/package/gentrace/providers/getters.py
@@ -1,4 +1,6 @@
 import os
+import re
+from urllib.parse import urlparse
 
 import openai
 
@@ -7,6 +9,9 @@
 openai.api_key = os.getenv("OPENAI_KEY")
 
 
+VALID_GENTRACE_HOST = r"^https?://[\w.-]+:\d{1,5}/api/v1/?$"
+
+
 def configure_openai():
     from gentrace import api_key, host
 
@@ -15,6 +20,9 @@ def configure_openai():
     if not api_key:
         raise ValueError("Gentrace API key not set")
 
+    if host and not re.match(VALID_GENTRACE_HOST, host):
+        raise ValueError("Gentrace host is invalid")
+
     gentrace_config = GentraceConfiguration(host=host)
     gentrace_config.access_token = api_key
 
diff --git a/package/tests/test_openai_embedding.py b/package/tests/test_openai_embedding.py
index 35c2e7d..71c3628 100644
--- a/package/tests/test_openai_embedding.py
+++ b/package/tests/test_openai_embedding.py
@@ -124,10 +124,12 @@ def test_openai_embedding_self_contained_pipeline_id_server(
     result = openai.Embedding.create(
         input="sample text",
         model="text-similarity-davinci-001",
-        pipeline_id="testing-value",
+        pipeline_id="testing-value-vivek",
     )
 
     assert uuid.UUID(result.pipeline_run_id) is not None
+
+    print("pipeline_id: ", result.pipeline_run_id)
     print(setup_teardown_openai)
 
 
diff --git a/package/tests/test_usage.py b/package/tests/test_usage.py
index 04b7ff6..af5567d 100644
--- a/package/tests/test_usage.py
+++ b/package/tests/test_usage.py
@@ -2,11 +2,25 @@
 from unittest.mock import create_autospec
 
 import pytest
-from urllib3.response import HTTPResponse
 
 import gentrace
 
 
+def test_gentrace_host_valid():
+    gentrace.api_key = os.getenv("GENTRACE_API_KEY")
+    gentrace.host = "http://localhost:3000/"
+
+    with pytest.raises(ValueError):
+        gentrace.configure_openai()
+
+    gentrace.host = "http://localhost:3000/api/v1/feedback"
+    with pytest.raises(ValueError):
+        gentrace.configure_openai()
+
+    gentrace.host = ""
+    gentrace.api_key = ""
+
+
 def test_openai_configure_should_raise_error():
     with pytest.raises(ValueError):
         gentrace.configure_openai()