1 files changed, 104 insertions, 0 deletions
diff --git a/ai.py b/ai.py
new file mode 100644
index 0000000..dd94cf8
--- /dev/null
+++ b/ai.py
@@ -0,0 +1,104 @@
+import aiohttp
+import asyncio
+import json
+import gconfig
+
+try:
+    from config import *
+except:
+    gconfig.copy_config()
+finally:
+    from config import *
+
+async def make_request(url, headers, payload):
+    async with aiohttp.ClientSession(headers=headers) as session:
+        async with session.post(url, json=payload) as response:
+            if payload.get("stream") == True:
+                async for line in response.content:
+                    line = line.decode("utf-8").strip()
+                    if line.startswith("data:"):
+                        data = line[5:].strip()
+                        if data == "DONE":
+                            break
+                        else:
+                            # data = await response.text()
+                            yield data
+            else:
+                response_data = b""
+                async for chunk in response.content.iter_chunks():
+                    response_data += chunk[0]
+                yield response_data.decode()
+
+async def completions(
+    url, 
+    api_key, 
+    model, 
+    messages, 
+    stream=False, 
+    max_tokens=MAX_TOKENS, 
+    extra=None
+):
+    print(url)
+
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    
+    payload = {
+        "messages": messages,
+        "model": model,
+        "stream": stream,
+        "max_tokens": max_tokens,
+    }
+    
+    if extra:
+        payload = payload | extra
+    
+    if stream:
+        async for data in make_request(url+"/chat/completions", headers, payload):
+            try:
+                yield json.loads(data)
+            except:
+                yield data
+    else:
+        async for data in make_request(url+"/chat/completions", headers, payload):
+            try:
+                yield json.loads(data)
+            except:
+                yield data
+
+async def responses(
+    url, 
+    api_key, 
+    model, 
+    input, 
+    stream=False, 
+    max_tokens=MAX_TOKENS, 
+    extra=None
+):
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "model": model,
+        "input": input,
+        "stream": stream,
+        "max_tokens": max_tokens,
+    }
+    if extra:
+        payload = payload | extra
+    
+    if stream:
+        async for data in make_request(url+"/chat/completions", headers, payload):
+            try:
+                yield json.loads(data)
+            except:
+                yield data
+    else:
+        async for data in make_request(url+"/chat/completions", headers, payload):
+            try:
+                yield json.loads(data)
+            except:
+                yield data