summaryrefslogtreecommitdiffstats
path: root/g4f/stubs.py
diff options
context:
space:
mode:
authorH Lohaus <hlohaus@users.noreply.github.com>2024-02-23 03:08:30 +0100
committerGitHub <noreply@github.com>2024-02-23 03:08:30 +0100
commitcf87d467c927d4b93916a42e7139f4e801172d62 (patch)
tree175b39f8486ce72be00a40cc345834abe859c84c /g4f/stubs.py
parent~ (diff)
parentFix unittests, use Union typing (diff)
downloadgpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.tar
gpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.tar.gz
gpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.tar.bz2
gpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.tar.lz
gpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.tar.xz
gpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.tar.zst
gpt4free-cf87d467c927d4b93916a42e7139f4e801172d62.zip
Diffstat (limited to 'g4f/stubs.py')
-rw-r--r--g4f/stubs.py89
1 files changed, 76 insertions, 13 deletions
diff --git a/g4f/stubs.py b/g4f/stubs.py
index 1cbbb134..49cf8a88 100644
--- a/g4f/stubs.py
+++ b/g4f/stubs.py
@@ -1,35 +1,98 @@
from __future__ import annotations
+from typing import Union
+
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: Union[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
+ content: Union[str, None] = None
+
+ def __init__(self, content: Union[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: Union[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