3 files changed, 159 insertions, 162 deletions

diff --git a/g4f/api/__init__.py b/g4f/api/__init__.py
index 3f0778a1..9033aafe 100644
--- a/g4f/api/__init__.py
+++ b/g4f/api/__init__.py
@@ -1,21 +1,27 @@
-import ast
 import logging
-import time
 import json
-import random
-import string
 import uvicorn
 import nest_asyncio
 
 from fastapi           import FastAPI, Response, Request
-from fastapi.responses import StreamingResponse
-from typing            import List, Union, Any, Dict, AnyStr
-#from ._tokenizer       import tokenize
+from fastapi.responses import StreamingResponse, RedirectResponse, HTMLResponse, JSONResponse
+from pydantic          import BaseModel
+from typing            import List
 
 import g4f
-from .. import debug
-
-debug.logging = True
+import g4f.debug
+from g4f.client import Client
+from g4f.typing import Messages
+
+class ChatCompletionsConfig(BaseModel):
+    messages: Messages
+    model: str
+    provider: str | None 
+    stream: bool = False
+    temperature: float | None 
+    max_tokens: int = None
+    stop: list[str] | str | None
+    access_token: str | None
 
 class Api:
     def __init__(self, engine: g4f, debug: bool = True, sentry: bool = False,
@@ -25,169 +31,82 @@ class Api:
         self.sentry = sentry
         self.list_ignored_providers = list_ignored_providers
 
-        self.app = FastAPI()
+        if debug:
+            g4f.debug.logging = True
+        self.client = Client()
+
         nest_asyncio.apply()
+        self.app = FastAPI()
 
-        JSONObject = Dict[AnyStr, Any]
-        JSONArray = List[Any]
-        JSONStructure = Union[JSONArray, JSONObject]
+        self.routes()
 
+    def routes(self):
         @self.app.get("/")
         async def read_root():
-            return Response(content=json.dumps({"info": "g4f API"}, indent=4), media_type="application/json")
+            return RedirectResponse("/v1", 302)
 
         @self.app.get("/v1")
         async def read_root_v1():
-            return Response(content=json.dumps({"info": "Go to /v1/chat/completions or /v1/models."}, indent=4), media_type="application/json")
+            return HTMLResponse('g4f API: Go to '
+                                '<a href="/v1/chat/completions">chat/completions</a> '
+                                'or <a href="/v1/models">models</a>.')
 
         @self.app.get("/v1/models")
         async def models():
-            model_list = []
-            for model in g4f.Model.__all__():
-                model_info = (g4f.ModelUtils.convert[model])
-                model_list.append({
-                'id': model,
+            model_list = dict(
+                (model, g4f.ModelUtils.convert[model])
+                for model in g4f.Model.__all__()
+            )
+            model_list = [{
+                'id': model_id,
                 'object': 'model',
                 'created': 0,
-                'owned_by': model_info.base_provider}
-                )
-            return Response(content=json.dumps({
-                'object': 'list',
-                'data': model_list}, indent=4), media_type="application/json")
+                'owned_by': model.base_provider
+            } for model_id, model in model_list.items()]
+            return JSONResponse(model_list)
 
         @self.app.get("/v1/models/{model_name}")
         async def model_info(model_name: str):
             try:
-                model_info = (g4f.ModelUtils.convert[model_name])
-
-                return Response(content=json.dumps({
+                model_info = g4f.ModelUtils.convert[model_name]
+                return JSONResponse({
                     'id': model_name,
                     'object': 'model',
                     'created': 0,
                     'owned_by': model_info.base_provider
-                }, indent=4), media_type="application/json")
+                })
             except:
-                return Response(content=json.dumps({"error": "The model does not exist."}, indent=4), media_type="application/json")
+                return JSONResponse({"error": "The model does not exist."})
 
         @self.app.post("/v1/chat/completions")
-        async def chat_completions(request: Request, item: JSONStructure = None):
-            item_data = {
-                'model': 'gpt-3.5-turbo',
-                'stream': False,
-            }
-
-            # item contains byte keys, and dict.get suppresses error
-            item_data.update({
-                key.decode('utf-8') if isinstance(key, bytes) else key: str(value)
-                for key, value in (item or {}).items()
-            })
-            # messages is str, need dict
-            if isinstance(item_data.get('messages'), str):
-                item_data['messages'] = ast.literal_eval(item_data.get('messages'))
-
-            model = item_data.get('model')
-            stream = True if item_data.get("stream") == "True" else False
-            messages = item_data.get('messages')
-            provider = item_data.get('provider', '').replace('g4f.Provider.', '')
-            provider = provider if provider and provider != "Auto" else None
-            temperature = item_data.get('temperature')
-
+        async def chat_completions(config: ChatCompletionsConfig = None, request: Request = None, provider: str = None):
             try:
-                response = g4f.ChatCompletion.create(
-                    model=model,
-                    stream=stream,
-                    messages=messages,
-                    temperature = temperature,
-                    provider = provider,
+                config.provider = provider if config.provider is None else config.provider
+                if config.access_token is None and request is not None:
+                    auth_header = request.headers.get("Authorization")
+                    if auth_header is not None:
+                        config.access_token = auth_header.split(None, 1)[-1]
+
+                response = self.client.chat.completions.create(
+                    **dict(config),
                     ignored=self.list_ignored_providers
                 )
             except Exception as e:
                 logging.exception(e)
-                content = json.dumps({
-                    "error": {"message": f"An error occurred while generating the response:\n{e}"},
-                    "model": model,
-                    "provider": g4f.get_last_provider(True)
-                })
-                return Response(content=content, status_code=500, media_type="application/json")
-            completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
-            completion_timestamp = int(time.time())
-
-            if not stream:
-                #prompt_tokens, _ = tokenize(''.join([message['content'] for message in messages]))
-                #completion_tokens, _ = tokenize(response)
-
-                json_data = {
-                    'id': f'chatcmpl-{completion_id}',
-                    'object': 'chat.completion',
-                    'created': completion_timestamp,
-                    'model': model,
-                    'provider': g4f.get_last_provider(True),
-                    'choices': [
-                        {
-                            'index': 0,
-                            'message': {
-                                'role': 'assistant',
-                                'content': response,
-                            },
-                            'finish_reason': 'stop',
-                        }
-                    ],
-                    'usage': {
-                        'prompt_tokens': 0, #prompt_tokens,
-                        'completion_tokens': 0, #completion_tokens,
-                        'total_tokens': 0, #prompt_tokens + completion_tokens,
-                    },
-                }
-
-                return Response(content=json.dumps(json_data, indent=4), media_type="application/json")
+                return Response(content=format_exception(e, config), status_code=500, media_type="application/json")
+
+            if not config.stream:
+                return JSONResponse(response.to_json())
 
             def streaming():
                 try:
                     for chunk in response:
-                        completion_data = {
-                            'id': f'chatcmpl-{completion_id}',
-                            'object': 'chat.completion.chunk',
-                            'created': completion_timestamp,
-                            'model': model,
-                            'provider': g4f.get_last_provider(True),
-                            'choices': [
-                                {
-                                    'index': 0,
-                                    'delta': {
-                                        'role': 'assistant',
-                                        'content': chunk,
-                                    },
-                                    'finish_reason': None,
-                                }
-                            ],
-                        }
-                        yield f'data: {json.dumps(completion_data)}\n\n'
-                        time.sleep(0.03)
-                    end_completion_data = {
-                        'id': f'chatcmpl-{completion_id}',
-                        'object': 'chat.completion.chunk',
-                        'created': completion_timestamp,
-                        'model': model,
-                        'provider': g4f.get_last_provider(True),
-                        'choices': [
-                            {
-                                'index': 0,
-                                'delta': {},
-                                'finish_reason': 'stop',
-                            }
-                        ],
-                    }
-                    yield f'data: {json.dumps(end_completion_data)}\n\n'
+                        yield f"data: {json.dumps(chunk.to_json())}\n\n"
                 except GeneratorExit:
                     pass
                 except Exception as e:
                     logging.exception(e)
-                    content = json.dumps({
-                        "error": {"message": f"An error occurred while generating the response:\n{e}"},
-                        "model": model,
-                        "provider": g4f.get_last_provider(True),
-                    })
-                    yield f'data: {content}'
+                    yield f'data: {format_exception(e, config)}'
 
             return StreamingResponse(streaming(), media_type="text/event-stream")
 
@@ -198,3 +117,11 @@ class Api:
     def run(self, ip):
         split_ip = ip.split(":")
         uvicorn.run(app=self.app, host=split_ip[0], port=int(split_ip[1]), use_colors=False)
+
+def format_exception(e: Exception, config: ChatCompletionsConfig) -> str:
+    last_provider = g4f.get_last_provider(True)
+    return json.dumps({
+        "error": {"message": f"ChatCompletionsError: {e.__class__.__name__}: {e}"},
+        "model": last_provider.get("model") if last_provider else config.model,
+        "provider": last_provider.get("name") if last_provider else config.provider
+    })
\ No newline at end of file
diff --git a/g4f/client.py b/g4f/client.py
index 4e5394b7..b44a5230 100644
--- a/g4f/client.py
+++ b/g4f/client.py
@@ -2,6 +2,9 @@ from __future__ import annotations
 
 import re
 import os
+import time
+import random
+import string
 
 from .stubs import ChatCompletion, ChatCompletionChunk, Image, ImagesResponse
 from .typing import Union, Generator, Messages, ImageType
@@ -10,10 +13,11 @@ from .image import ImageResponse as ImageProviderResponse
 from .Provider.BingCreateImages import BingCreateImages
 from .Provider.needs_auth import Gemini, OpenaiChat
 from .errors import NoImageResponseError
-from . import get_model_and_provider
+from . import get_model_and_provider, get_last_provider
 
 ImageProvider = Union[BaseProvider, object]
 Proxies = Union[dict, str]
+IterResponse = Generator[ChatCompletion | ChatCompletionChunk, None, None]
 
 def read_json(text: str) -> dict:
     """
@@ -31,18 +35,16 @@ def read_json(text: str) -> dict:
     return text
 
 def iter_response(
-    response: iter,
+    response: iter[str],
     stream: bool,
     response_format: dict = None,
     max_tokens: int = None,
     stop: list = None
-) -> Generator:
+) -> IterResponse:
     content = ""
     finish_reason = None
-    last_chunk = None
+    completion_id = ''.join(random.choices(string.ascii_letters + string.digits, k=28))
     for idx, chunk in enumerate(response):
-        if last_chunk is not None:
-            yield ChatCompletionChunk(last_chunk, finish_reason)
         content += str(chunk)
         if max_tokens is not None and idx + 1 >= max_tokens:
             finish_reason = "length"
@@ -63,16 +65,25 @@ def iter_response(
         if first != -1:
             finish_reason = "stop"
         if stream:
-            last_chunk = chunk
+            yield ChatCompletionChunk(chunk, None, completion_id, int(time.time()))
         if finish_reason is not None:
             break
-    if last_chunk is not None:
-        yield ChatCompletionChunk(last_chunk, finish_reason)
-    if not stream:
+    finish_reason = "stop" if finish_reason is None else finish_reason
+    if stream:
+        yield ChatCompletionChunk(None, finish_reason, completion_id, int(time.time()))
+    else:
         if response_format is not None and "type" in response_format:
             if response_format["type"] == "json_object":
                 content = read_json(content)
-        yield ChatCompletion(content, finish_reason)
+        yield ChatCompletion(content, finish_reason, completion_id, int(time.time()))
+
+def iter_append_model_and_provider(response: IterResponse) -> IterResponse:
+    last_provider = None
+    for chunk in response:
+        last_provider = get_last_provider(True) if last_provider is None else last_provider
+        chunk.model = last_provider.get("model")
+        chunk.provider =  last_provider.get("name")
+        yield chunk
 
 class Client():
     proxies: Proxies = None
@@ -113,7 +124,7 @@ class Completions():
         stream: bool = False,
         response_format: dict = None,
         max_tokens: int = None,
-        stop: Union[list. str] = None,
+        stop: list[str] | str = None,
         **kwargs
     ) -> Union[ChatCompletion, Generator[ChatCompletionChunk]]:
         if max_tokens is not None:
@@ -128,7 +139,7 @@ class Completions():
         )
         response = provider.create_completion(model, messages, stream=stream, proxy=self.client.get_proxy(), **kwargs)
         stop = [stop] if isinstance(stop, str) else stop
-        response = iter_response(response, stream, response_format, max_tokens, stop)
+        response = iter_append_model_and_provider(iter_response(response, stream, response_format, max_tokens, stop))
         return response if stream else next(response)
 
 class Chat():
diff --git a/g4f/stubs.py b/g4f/stubs.py
index 1cbbb134..b9934b8c 100644
--- a/g4f/stubs.py
+++ b/g4f/stubs.py
@@ -2,34 +2,93 @@
 from __future__ import annotations
 
 class Model():
-    def __getitem__(self, item):
-        return getattr(self, item)
+    ...
 
 class ChatCompletion(Model):
-    def __init__(self, content: str, finish_reason: str):
-        self.choices = [ChatCompletionChoice(ChatCompletionMessage(content, finish_reason))]
+    def __init__(
+        self,
+        content: str,
+        finish_reason: str,
+        completion_id: str = None,
+        created: int = None
+    ):
+        self.id: str = f"chatcmpl-{completion_id}" if completion_id else None
+        self.object: str = "chat.completion"
+        self.created: int = created
+        self.model: str = None
+        self.provider: str = None
+        self.choices = [ChatCompletionChoice(ChatCompletionMessage(content), finish_reason)]
+        self.usage: dict[str, int] = {
+            "prompt_tokens": 0, #prompt_tokens,
+            "completion_tokens": 0, #completion_tokens,
+            "total_tokens": 0, #prompt_tokens + completion_tokens,
+        }
+
+    def to_json(self):
+        return {
+            **self.__dict__,
+            "choices": [choice.to_json() for choice in self.choices]
+        }
 
 class ChatCompletionChunk(Model):
-    def __init__(self, content: str, finish_reason: str):
-        self.choices = [ChatCompletionDeltaChoice(ChatCompletionDelta(content, finish_reason))]
+    def __init__(
+        self,
+        content: str,
+        finish_reason: str,
+        completion_id: str = None,
+        created: int = None
+    ):
+        self.id: str = f"chatcmpl-{completion_id}" if completion_id else None
+        self.object: str = "chat.completion.chunk"
+        self.created: int = created
+        self.model: str = None
+        self.provider: str = None
+        self.choices = [ChatCompletionDeltaChoice(ChatCompletionDelta(content), finish_reason)]
+
+    def to_json(self):
+        return {
+            **self.__dict__,
+            "choices": [choice.to_json() for choice in self.choices]
+        }
 
 class ChatCompletionMessage(Model):
-    def __init__(self, content: str, finish_reason: str):
+    def __init__(self, content: str | None):
+        self.role = "assistant"
         self.content = content
-        self.finish_reason = finish_reason
+
+    def to_json(self):
+        return self.__dict__
 
 class ChatCompletionChoice(Model):
-    def __init__(self, message: ChatCompletionMessage):
+    def __init__(self, message: ChatCompletionMessage, finish_reason: str):
+        self.index = 0
         self.message = message
+        self.finish_reason = finish_reason
+
+    def to_json(self):
+        return {
+            **self.__dict__,
+            "message": self.message.to_json()
+        }
 
 class ChatCompletionDelta(Model):
-    def __init__(self, content: str, finish_reason: str):
-        self.content = content
-        self.finish_reason = finish_reason
+    def __init__(self, content: str | None):
+        if content is not None:
+            self.content = content
+
+    def to_json(self):
+        return self.__dict__
 
 class ChatCompletionDeltaChoice(Model):
-    def __init__(self, delta: ChatCompletionDelta):
+    def __init__(self, delta: ChatCompletionDelta, finish_reason: str | None):
         self.delta = delta
+        self.finish_reason = finish_reason
+
+    def to_json(self):
+        return {
+            **self.__dict__,
+            "delta": self.delta.to_json()
+        }
 
 class Image(Model):
     url: str