use Messages API (OpenAI)

rag.py

@@ -35,27 +35,28 @@ class RAG:
         
     def predict(self, instruction, context, model_parameters):
 
-        api_key = os.getenv("HF_TOKEN")
-
-
-        headers = {
-        "Accept" : "application/json",
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json" 
-        }
-
-        query = f"### Instruction\n{instruction}\n\n### Context\n{context}\n\n### Answer\n "
-        #prompt = "You are a helpful assistant. Answer the question using only the context you are provided with. If it is not possible to do it with the context, just say 'I can't answer'. <|endoftext|>"
-
-
-        payload = {
-        "inputs": query,
-        "parameters": model_parameters
-        }
+        from openai import OpenAI
         
-        response = requests.post(self.model_name, headers=headers, json=payload)
+        # init the client but point it to TGI
+        client = OpenAI(
+            base_url=os.getenv("MODEL")+ "/v1/",
+            api_key=os.getenv("HF_TOKEN")
+        )
+
+        query = f"{context}\n\n{instruction}"
+        #sys_prompt = "You are a helpful assistant. Answer the question using only the context you are provided with. If it is not possible to do it with the context, just say 'I can't answer'. <|endoftext|>"
+
+        chat_completion = client.chat.completions.create(
+            model="tgi",
+            messages=[
+                #{"role": "system", "content": sys_prompt },
+                {"role": "user", "content": query}
+            ],
+            stream=False
+        )
+
+        return(chat_completion)
 
-        return response.json()[0]["generated_text"].split("###")[-1][8:]
     
     def beautiful_context(self, docs):