Merge pull request #81 from SkywardAI/add_rag_completion_test_case

Update test_apis.py to implement rag-completion test
SkywardAI · Sep 3, 2024 · 1b2207f · 1b2207f
2 parents cea1794 + 52bbcd0
commit 1b2207f
Showing 1 changed file with 27 additions and 1 deletion.
diff --git a/example/test_apis.py b/example/test_apis.py
@@ -82,6 +82,32 @@ def test_inference_by_request_stream(self):
         }
         res=requests.post(url=self.base_url+"/v1/chat/completions", json=data,headers={"Content-Type": "application/json", "Authorization":"Bearer no-key"},)
         self.assertEqual(res.status_code, 200)
+
+    def test_rag_completion(self):
+        #Load the dataset
+        dataset = {
+                    "name": "aisuko/squad01-v2",
+                    "url": "https://datasets-server.huggingface.co/rows?dataset=aisuko%2Fsquad01-v2&config=default&split=validation&offset=0&length=100"
+                }
+        loadDataset = requests.post(url=self.base_url+"/v1/embeddings/dataset", json=dataset, headers={"Content-Type": "application/json", "Authorization":"Bearer no-key"})
+        self.assertEqual(loadDataset.status_code, 200)
+
+        #Test RAG completions
+        data = {
+                "messages": [
+                    {
+                    "role": "system",
+                    "content": "You are a helpful assistant who helps users solve their questions."
+                    },
+                    {
+                    "role": "user",
+                    "content": "tell me something interest about massachusetts"
+                    }
+                ],
+                "dataset_name": "aisuko/squad01-v2"
+            }
+        res = requests.post(url=self.base_url+"/v1/chat/rag-completions", json=data, headers={"Content-Type": "application/json", "Authorization":"Bearer no-key"})
+        self.assertEqual(res.status_code, 200)
 
     def test_embedding_api(self):
         data = {
@@ -101,4 +127,4 @@ def test_version_api(self):
         res = requests.get(url=self.base_url+"/v1/version")
         self.assertEqual(res.status_code, 200)
         data = res.json().get("inference_engine_version")
-        self.assertEqual('server--b1-2321a5e', data)
+        self.assertEqual('server--b1-27d4b7c', data)