0.3.4.0

1 files changed, 148 insertions, 36 deletions

diff --git a/g4f/Provider/Airforce.py b/g4f/Provider/Airforce.py
index c7ae44c0..6254e160 100644
--- a/g4f/Provider/Airforce.py
+++ b/g4f/Provider/Airforce.py
@@ -1,59 +1,171 @@
 from __future__ import annotations
-from typing import Any, Dict
-import inspect
 
-from aiohttp import ClientSession
+import random
+import json
+import re
 
 from ..typing import AsyncResult, Messages
 from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
-from .helper import format_prompt
+from ..image import ImageResponse
+from ..requests import StreamSession, raise_for_status
 from .airforce.AirforceChat import AirforceChat
 from .airforce.AirforceImage import AirforceImage
 
 class Airforce(AsyncGeneratorProvider, ProviderModelMixin):
     url = "https://api.airforce"
     api_endpoint_completions = AirforceChat.api_endpoint
-    api_endpoint_imagine2 = AirforceImage.api_endpoint
+    api_endpoint_imagine = AirforceImage.api_endpoint
     working = True
-    supports_stream = AirforceChat.supports_stream
-    supports_system_message = AirforceChat.supports_system_message
-    supports_message_history = AirforceChat.supports_message_history 
-    
-    default_model = AirforceChat.default_model
-    models = [*AirforceChat.models, *AirforceImage.models]
-    
+    default_model = "gpt-4o-mini"
+    supports_system_message = True
+    supports_message_history = True
+    text_models = [
+        'gpt-4-turbo',
+        default_model,
+        'llama-3.1-70b-turbo',
+        'llama-3.1-8b-turbo',
+    ]
+    image_models = [
+        'flux',
+        'flux-realism',
+        'flux-anime',
+        'flux-3d',
+        'flux-disney',
+        'flux-pixel',
+        'flux-4o',
+        'any-dark',
+    ]
+    models = [
+        *text_models,
+        *image_models,
+    ]
     model_aliases = {
-        **AirforceChat.model_aliases,
-        **AirforceImage.model_aliases
+        "gpt-4o": "chatgpt-4o-latest",
+        "llama-3.1-70b": "llama-3.1-70b-turbo",
+        "llama-3.1-8b": "llama-3.1-8b-turbo",
+        "gpt-4": "gpt-4-turbo",
     }
 
     @classmethod
-    def get_model(cls, model: str) -> str:
-        if model in cls.models:
-            return model
-        elif model in cls.model_aliases:
-            return cls.model_aliases[model]
+    def create_async_generator(
+        cls,
+        model: str,
+        messages: Messages,
+        proxy: str = None,
+        seed: int = None,
+        size: str = "1:1",
+        stream: bool = False,
+        **kwargs
+    ) -> AsyncResult:
+        model = cls.get_model(model)
+
+        if model in cls.image_models:
+            return cls._generate_image(model, messages, proxy, seed, size)
         else:
-            return cls.default_model
+            return cls._generate_text(model, messages, proxy, stream, **kwargs)
 
     @classmethod
-    async def create_async_generator(cls, model: str, messages: Messages, **kwargs) -> AsyncResult:
-        model = cls.get_model(model)
-        
-        provider = AirforceChat if model in AirforceChat.text_models else AirforceImage
+    async def _generate_image(
+        cls,
+        model: str,
+        messages: Messages,
+        proxy: str = None,
+        seed: int = None,
+        size: str = "1:1",
+        **kwargs
+    ) -> AsyncResult:
+        headers = {
+            "accept": "*/*",
+            "accept-language": "en-US,en;q=0.9",
+            "cache-control": "no-cache",
+            "origin": "https://llmplayground.net",
+            "user-agent": "Mozilla/5.0"
+        }
+        if seed is None:
+            seed = random.randint(0, 100000)
+        prompt = messages[-1]['content']
 
-        if model not in provider.models:
-            raise ValueError(f"Unsupported model: {model}")
+        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()
 
-        # Get the signature of the provider's create_async_generator method
-        sig = inspect.signature(provider.create_async_generator)
-        
-        # Filter kwargs to only include parameters that the provider's method accepts
-        filtered_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
+                if 'application/json' in content_type:
+                    raise RuntimeError(await response.json().get("error", {}).get("message"))
+                elif 'image' in content_type:
+                    image_data = b""
+                    async for chunk in response.iter_content():
+                        if chunk:
+                            image_data += chunk
+                    image_url = f"{cls.api_endpoint_imagine}?model={model}&prompt={prompt}&size={size}&seed={seed}"
+                    yield ImageResponse(images=image_url, alt=prompt)
 
-        # Add model and messages to filtered_kwargs
-        filtered_kwargs['model'] = model
-        filtered_kwargs['messages'] = messages
+    @classmethod
+    async def _generate_text(
+        cls,
+        model: str,
+        messages: Messages,
+        proxy: str = None,
+        stream: bool = False,
+        max_tokens: int = 4096,
+        temperature: float = 1,
+        top_p: float = 1,
+        **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"
+        }
+        async with StreamSession(headers=headers, proxy=proxy) as session:
+            data = {
+                "messages": messages,
+                "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:
+                    json_data = await response.json()
+                    content = json_data['choices'][0]['message']['content']
+                    yield cls._filter_content(content)
 
-        async for result in provider.create_async_generator(**filtered_kwargs):
-            yield result
+    @classmethod
+    def _filter_content(cls, part_response: str) -> str:
+        part_response = re.sub(
+            r"One message exceeds the \d+chars per message limit\..+https:\/\/discord\.com\/invite\/\S+",
+            '',
+            part_response
+        )
+        
+        part_response = re.sub(
+            r"Rate limit \(\d+\/minute\) exceeded\. Join our discord for more: .+https:\/\/discord\.com\/invite\/\S+",
+            '',
+            part_response
+        )
+        return part_response
\ No newline at end of file