File size: 731 Bytes
b93b2dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import requests
from dotenv import load_dotenv

load_dotenv()

class LLM:
    def __init__(self, url, api_key):
        self.endpoint = url
        self.api_key = api_key

    def generate_response(self, prompt):
        headers = {
            "Content-Type": "application/json",
            "api-key": self.api_key,
        }

        data = {
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 1500,
            "temperature": 0.5,
        }

        response = requests.post(self.endpoint, headers=headers, json=data)

        if response.status_code == 200:
            return response.json()["choices"][0]["message"]["content"]
        else:
            return ValueError(response.text)