From 74397096b794631e718e7e5dfc7ed8517d0e42c2 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Fri, 23 Feb 2024 02:35:13 +0100
Subject: Use new client in inter api

---
 g4f/stubs.py | 85 ++++++++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 72 insertions(+), 13 deletions(-)

(limited to 'g4f/stubs.py')

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
-- 
cgit v1.2.3