ARGO.py

#
# A wrapper class for the Argonne Argo LLM service
#

import os
import requests
import json

MODEL_GPT35 = "gpt35"
MODEL_GPT4 = "gpt4"

class ArgoWrapper:
    default_url = "https://apps-dev.inside.anl.gov/argoapi/api/v1/resource/chat/"

    def __init__(self, 
                 url = None, 
                 model = MODEL_GPT35, 
                 system = "",
                 temperature = 0.8, 
                 top_p=0.7, 
                 user = os.getenv("USER"))-> None:
        self.url = url
        if self.url is None:
            self.url = ArgoWrapper.default_url
        self.model = model
        self.temperature = temperature
        self.top_p = top_p
        self.user = user
        self.system = ""

    def invoke(self, prompt: str):
        headers = {
            "Content-Type": "application/json"
        }
        data = {
                "user": self.user,
                "model": self.model,
                "system": self.system,
                "prompt": [prompt],
                "stop": [],
                "temperature": self.temperature,
                "top_p": self.top_p
        }
            
        data_json = json.dumps(data)    
        response = requests.post(self.url, headers=headers, data=data_json)

        if response.status_code == 200:
            parsed = json.loads(response.text)
            return parsed
        else:
            raise Exception(f"Request failed with status code: {response.status_code}")