Skip to content

Commit

Permalink
refactor(Airforce): improve API handling and add authentication
Browse files Browse the repository at this point in the history 
- Implement API key authentication with check_api_key method
- Refactor image generation to use new imagine2 endpoint
- Improve text generation with better error handling and streaming
- Update model aliases and add new image models
- Enhance content filtering for various model outputs
- Replace StreamSession with aiohttp's ClientSession for async operations
- Simplify model fetching logic and remove redundant code
- Add is_image_model method for better model type checking
- Update class attributes for better organization and clarity
  • Loading branch information
kqlio67 committed Dec 2, 2024
1 parent f01b61c commit 569ee15
Showing 1 changed file with 162 additions and 159 deletions.
321 changes: 162 additions & 159 deletions g4f/Provider/Airforce.py
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,17 @@
from __future__ import annotations

import random
import json
import random
import re

import requests
from requests.packages.urllib3.exceptions import InsecureRequestWarning
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
from urllib.parse import quote
from aiohttp import ClientSession

from ..typing import AsyncResult, Messages
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ..image import ImageResponse
from ..requests import StreamSession, raise_for_status
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin

requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

def split_message(message: str, max_length: int = 1000) -> list[str]:
"""Splits the message into parts up to (max_length)."""
Expand All @@ -30,197 +29,198 @@ def split_message(message: str, max_length: int = 1000) -> list[str]:
class Airforce(AsyncGeneratorProvider, ProviderModelMixin):
url = "https://llmplayground.net"
api_endpoint_completions = "https://api.airforce/chat/completions"
api_endpoint_imagine = "https://api.airforce/imagine2"
api_endpoint_imagine2 = "https://api.airforce/v1/imagine2"

working = True
needs_auth = True
supports_stream = True
supports_system_message = True
supports_message_history = True

@classmethod
def fetch_completions_models(cls):
response = requests.get('https://api.airforce/models', verify=False)
response.raise_for_status()
data = response.json()
return [model['id'] for model in data['data']]

@classmethod
def fetch_imagine_models(cls):
response = requests.get('https://api.airforce/imagine/models', verify=False)
response.raise_for_status()
return response.json()

default_model = "gpt-4o-mini"
default_image_model = "flux"
additional_models_imagine = ["stable-diffusion-xl-base", "stable-diffusion-xl-lightning", "flux-1.1-pro"]

@classmethod
def get_models(cls):
if not cls.models:
cls.image_models = [*cls.fetch_imagine_models(), *cls.additional_models_imagine]
cls.models = [
*cls.fetch_completions_models(),
*cls.image_models
]
return cls.models

model_aliases = {
### completions ###
# openchat
additional_models_imagine = ["flux-1.1-pro", "dall-e-3"]

model_aliases = {
# Alias mappings for models
"openchat-3.5": "openchat-3.5-0106",

# deepseek-ai
"deepseek-coder": "deepseek-coder-6.7b-instruct",

# NousResearch
"hermes-2-dpo": "Nous-Hermes-2-Mixtral-8x7B-DPO",
"hermes-2-pro": "hermes-2-pro-mistral-7b",

# teknium
"openhermes-2.5": "openhermes-2.5-mistral-7b",

# liquid
"lfm-40b": "lfm-40b-moe",

# DiscoResearch
"german-7b": "discolm-german-7b-v1",

# meta-llama
"discolm-german-7b": "discolm-german-7b-v1",
"llama-2-7b": "llama-2-7b-chat-int8",
"llama-2-7b": "llama-2-7b-chat-fp16",
"llama-3.1-70b": "llama-3.1-70b-chat",
"llama-3.1-8b": "llama-3.1-8b-chat",
"llama-3.1-70b": "llama-3.1-70b-turbo",
"llama-3.1-8b": "llama-3.1-8b-turbo",

# inferless
"neural-7b": "neural-chat-7b-v3-1",

# HuggingFaceH4
"zephyr-7b": "zephyr-7b-beta",


### imagine ###
"sdxl": "stable-diffusion-xl-base",
"sdxl": "stable-diffusion-xl-lightning",
"flux-pro": "flux-1.1-pro",
}

@classmethod
def create_async_generator(
def fetch_completions_models(cls):
response = requests.get('https://api.airforce/models', verify=False)
response.raise_for_status()
data = response.json()
return [model['id'] for model in data['data']]

@classmethod
def fetch_imagine_models(cls):
response = requests.get(
'https://api.airforce/imagine/models',
'https://api.airforce/v1/imagine2/models',
verify=False
)
response.raise_for_status()
return response.json()

image_models = fetch_imagine_models.__get__(None, object)() + additional_models_imagine

@classmethod
def is_image_model(cls, model: str) -> bool:
return model in cls.image_models

models = list(dict.fromkeys([default_model] +
fetch_completions_models.__get__(None, object)() +
image_models))

@classmethod
async def check_api_key(cls, api_key: str) -> bool:
"""
Always returns True to allow all models when missing_key is used.
"""
if not api_key or api_key == "missing_key":
return True # No restrictions if no key or missing_key is used.

headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Accept": "*/*",
}

try:
async with ClientSession(headers=headers) as session:
async with session.get(f"https://api.airforce/check?key={api_key}") as response:
if response.status == 200:
data = await response.json()
return data.get('info') in ['Sponsor key', 'Premium key']
return False
except Exception as e:
print(f"Error checking API key: {str(e)}")
return False

@classmethod
async def generate_image(
cls,
model: str,
messages: Messages,
proxy: str = None,
prompt: str = None,
seed: int = None,
size: str = "1:1", # "1:1", "16:9", "9:16", "21:9", "9:21", "1:2", "2:1"
stream: bool = False,
**kwargs
prompt: str,
api_key: str,
size: str,
seed: int,
proxy: str = None
) -> AsyncResult:
model = cls.get_model(model)
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:132.0) Gecko/20100101 Firefox/132.0",
"Authorization": f"Bearer {api_key}",
}
params = {"model": model, "prompt": prompt, "size": size, "seed": seed}

if model in cls.image_models:
if prompt is None:
prompt = messages[-1]['content']
return cls._generate_image(model, prompt, proxy, seed, size)
else:
return cls._generate_text(model, messages, proxy, stream, **kwargs)
async with ClientSession(headers=headers) as session:
async with session.get(cls.api_endpoint_imagine2, params=params, proxy=proxy) as response:
if response.status == 200:
image_url = str(response.url)
yield ImageResponse(images=image_url, alt=f"Generated image: {prompt}")
else:
error_text = await response.text()
raise RuntimeError(f"Image generation failed: {response.status} - {error_text}")

@classmethod
async def _generate_image(
async def generate_text(
cls,
model: str,
prompt: str,
proxy: str = None,
seed: int = None,
size: str = "1:1",
**kwargs
messages: Messages,
max_tokens: int,
temperature: float,
top_p: float,
stream: bool,
api_key: str,
proxy: str = None
) -> AsyncResult:
headers = {
"accept": "*/*",
"accept-language": "en-US,en;q=0.9",
"cache-control": "no-cache",
"user-agent": "Mozilla/5.0"
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
full_message = "\n".join([msg['content'] for msg in messages])
message_chunks = split_message(full_message, max_length=1000)

data = {
"messages": [{"role": "user", "content": chunk} for chunk in message_chunks],
"model": model,
"max_tokens": max_tokens,
"temperature": temperature,
"top_p": top_p,
"stream": stream,
}
if seed is None:
seed = random.randint(0, 100000)

async with StreamSession(headers=headers, proxy=proxy) as session:
params = {
"model": model,
"prompt": prompt,
"size": size,
"seed": seed
}
async with session.get(f"{cls.api_endpoint_imagine}", params=params) as response:
await raise_for_status(response)
content_type = response.headers.get('Content-Type', '').lower()

if 'application/json' in content_type:
raise RuntimeError(await response.json().get("error", {}).get("message"))
elif content_type.startswith("image/"):
image_url = f"{cls.api_endpoint_imagine}?model={model}&prompt={quote(prompt)}&size={size}&seed={seed}"
yield ImageResponse(images=image_url, alt=prompt)

async with ClientSession(headers=headers) as session:
async with session.post(cls.api_endpoint_completions, json=data, proxy=proxy) as response:
response.raise_for_status()

if stream:
async for line in response.content:
line = line.decode('utf-8').strip()
if line.startswith('data: '):
try:
json_str = line[6:] # Remove 'data: ' prefix
chunk = json.loads(json_str)
if 'choices' in chunk and chunk['choices']:
delta = chunk['choices'][0].get('delta', {})
if 'content' in delta:
filtered_content = cls._filter_response(delta['content'])
yield filtered_content
except json.JSONDecodeError:
continue
else:
# Non-streaming response
result = await response.json()
if 'choices' in result and result['choices']:
message = result['choices'][0].get('message', {})
content = message.get('content', '')
filtered_content = cls._filter_response(content)
yield filtered_content

@classmethod
async def _generate_text(
async def create_async_generator(
cls,
model: str,
messages: Messages,
prompt: str = None,
proxy: str = None,
stream: bool = False,
max_tokens: int = 4096,
temperature: float = 1,
top_p: float = 1,
stream: bool = True,
api_key: str = "missing_key",
size: str = "1:1",
seed: int = None,
**kwargs
) -> AsyncResult:
headers = {
"accept": "*/*",
"accept-language": "en-US,en;q=0.9",
"authorization": "Bearer missing api key",
"content-type": "application/json",
"user-agent": "Mozilla/5.0"
}
# Always proceed even if API key is "missing_key"
if not await cls.check_api_key(api_key):
print("Warning: Using invalid or missing API key. All models are still accessible.")

full_message = "\n".join(
[f"{msg['role'].capitalize()}: {msg['content']}" for msg in messages]
)

message_chunks = split_message(full_message, max_length=1000)

async with StreamSession(headers=headers, proxy=proxy) as session:
full_response = ""
for chunk in message_chunks:
data = {
"messages": [{"role": "user", "content": chunk}],
"model": model,
"max_tokens": max_tokens,
"temperature": temperature,
"top_p": top_p,
"stream": stream
}

async with session.post(cls.api_endpoint_completions, json=data) as response:
await raise_for_status(response)
content_type = response.headers.get('Content-Type', '').lower()

if 'application/json' in content_type:
json_data = await response.json()
if json_data.get("model") == "error":
raise RuntimeError(json_data['choices'][0]['message'].get('content', ''))
if stream:
async for line in response.iter_lines():
if line:
line = line.decode('utf-8').strip()
if line.startswith("data: ") and line != "data: [DONE]":
json_data = json.loads(line[6:])
content = json_data['choices'][0]['delta'].get('content', '')
if content:
yield cls._filter_content(content)
else:
content = json_data['choices'][0]['message']['content']
full_response += cls._filter_content(content)

yield full_response
if cls.is_image_model(model):
if prompt is None:
prompt = messages[-1]['content']
if seed is None:
seed = random.randint(0, 10000)

async for result in cls.generate_image(model, prompt, api_key, size, seed, proxy):
yield result
else:
async for result in cls.generate_text(model, messages, max_tokens, temperature, top_p, stream, api_key, proxy):
yield result

@classmethod
def _filter_content(cls, part_response: str) -> str:
Expand All @@ -229,16 +229,19 @@ def _filter_content(cls, part_response: str) -> str:
'',
part_response
)

part_response = re.sub(
r"Rate limit \(\d+\/minute\) exceeded\. Join our discord for more: .+https:\/\/discord\.com\/invite\/\S+",
'',
part_response
)

part_response = re.sub(
r"\[ERROR\] '\w{8}-\w{4}-\w{4}-\w{4}-\w{12}'", # any-uncensored
'',
part_response
)

return part_response

@classmethod
def _filter_response(cls, response: str) -> str:
filtered_response = re.sub(r"\[ERROR\] '\w{8}-\w{4}-\w{4}-\w{4}-\w{12}'", '', response) # any-uncensored
filtered_response = re.sub(r'<\|im_end\|>', '', response) # hermes-2-pro-mistral-7b
filtered_response = re.sub(r'</s>', '', response) # neural-chat-7b-v3-1
filtered_response = cls._filter_content(filtered_response)
return filtered_response

0 comments on commit 569ee15

Please sign in to comment.