1 files changed, 142 insertions, 0 deletions

diff --git a/g4f/Provider/DeepInfraChat.py b/g4f/Provider/DeepInfraChat.py
new file mode 100644
index 00000000..fa6dd2fe
--- /dev/null
+++ b/g4f/Provider/DeepInfraChat.py
@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from aiohttp import ClientSession
+import json
+
+from ..typing import AsyncResult, Messages, ImageType
+from ..image import to_data_uri
+from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
+from .helper import format_prompt
+
+
+class DeepInfraChat(AsyncGeneratorProvider, ProviderModelMixin):
+    url = "https://deepinfra.com/chat"
+    api_endpoint = "https://api.deepinfra.com/v1/openai/chat/completions"
+    working = True
+    supports_stream = True
+    supports_system_message = True
+    supports_message_history = True
+    
+    default_model = 'meta-llama/Meta-Llama-3.1-70B-Instruct'
+    models = [
+        'meta-llama/Meta-Llama-3.1-405B-Instruct',
+        'meta-llama/Meta-Llama-3.1-70B-Instruct',
+        'meta-llama/Meta-Llama-3.1-8B-Instruct',
+        'mistralai/Mixtral-8x22B-Instruct-v0.1',
+        'mistralai/Mixtral-8x7B-Instruct-v0.1',
+        'microsoft/WizardLM-2-8x22B',
+        'microsoft/WizardLM-2-7B',
+        'Qwen/Qwen2-72B-Instruct',
+        'microsoft/Phi-3-medium-4k-instruct',
+        'google/gemma-2-27b-it',
+        'openbmb/MiniCPM-Llama3-V-2_5', # Image upload is available
+        'mistralai/Mistral-7B-Instruct-v0.3',
+        'lizpreciatior/lzlv_70b_fp16_hf',
+        'openchat/openchat-3.6-8b',
+        'Phind/Phind-CodeLlama-34B-v2',
+        'cognitivecomputations/dolphin-2.9.1-llama-3-70b',
+    ]
+    model_aliases = {
+        "llama-3.1-405b": "meta-llama/Meta-Llama-3.1-405B-Instruct",
+        "llama-3.1-70b": "meta-llama/Meta-Llama-3.1-70B-Instruct",
+        "Llama-3.1-8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
+        "mixtral-8x22b": "mistralai/Mixtral-8x22B-Instruct-v0.1",
+        "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        "wizardlm-2-8x22b": "microsoft/WizardLM-2-8x22B",
+        "wizardlm-2-7b": "microsoft/WizardLM-2-7B",
+        "qwen-2-72b": "Qwen/Qwen2-72B-Instruct",
+        "phi-3-medium-4k": "microsoft/Phi-3-medium-4k-instruct",
+        "gemma-2b-27b": "google/gemma-2-27b-it",
+        "minicpm-llama-3-v2.5": "openbmb/MiniCPM-Llama3-V-2_5",
+        "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.3",
+        "lzlv_70b": "lizpreciatior/lzlv_70b_fp16_hf",
+        "openchat-3.6-8b": "openchat/openchat-3.6-8b",
+        "phind-codellama-34b-v2": "Phind/Phind-CodeLlama-34B-v2",
+        "dolphin-2.9.1-llama-3-70b": "cognitivecomputations/dolphin-2.9.1-llama-3-70b",
+    }
+
+
+    @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]
+        else:
+            return cls.default_model
+
+    @classmethod
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: Messages,
+        proxy: str = None,
+        image: ImageType = None,
+        image_name: str = None,
+        **kwargs
+    ) -> AsyncResult:
+        model = cls.get_model(model)
+        
+        headers = {
+            'Accept-Language': 'en-US,en;q=0.9',
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive',
+            'Content-Type': 'application/json',
+            'Origin': 'https://deepinfra.com',
+            'Pragma': 'no-cache',
+            'Referer': 'https://deepinfra.com/',
+            'Sec-Fetch-Dest': 'empty',
+            'Sec-Fetch-Mode': 'cors',
+            'Sec-Fetch-Site': 'same-site',
+            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
+            'X-Deepinfra-Source': 'web-embed',
+            'accept': 'text/event-stream',
+            'sec-ch-ua': '"Not;A=Brand";v="24", "Chromium";v="128"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"',
+        }
+        
+        async with ClientSession(headers=headers) as session:
+            prompt = format_prompt(messages)
+            data = {
+                'model': model,
+                'messages': [
+                    {'role': 'system', 'content': 'Be a helpful assistant'},
+                    {'role': 'user', 'content': prompt}
+                ],
+                'stream': True
+            }
+
+            if model == 'openbmb/MiniCPM-Llama3-V-2_5' and image is not None:
+                data['messages'][-1]['content'] = [
+                    {
+                        'type': 'image_url',
+                        'image_url': {
+                            'url': to_data_uri(image)
+                        }
+                    },
+                    {
+                        'type': 'text',
+                        'text': messages[-1]['content']
+                    }
+                ]
+
+            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
+                response.raise_for_status()
+                async for line in response.content:
+                    if line:
+                        decoded_line = line.decode('utf-8').strip()
+                        if decoded_line.startswith('data:'):
+                            json_part = decoded_line[5:].strip()
+                            if json_part == '[DONE]':
+                                break
+                            try:
+                                data = json.loads(json_part)
+                                choices = data.get('choices', [])
+                                if choices:
+                                    delta = choices[0].get('delta', {})
+                                    content = delta.get('content', '')
+                                    if content:
+                                        yield content
+                            except json.JSONDecodeError:
+                                print(f"JSON decode error: {json_part}")