stubs.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. from __future__ import annotations
  2. from typing import Union
  3. class Model():
  4. ...
  5. class ChatCompletion(Model):
  6. def __init__(
  7. self,
  8. content: str,
  9. finish_reason: str,
  10. completion_id: str = None,
  11. created: int = None
  12. ):
  13. self.id: str = f"chatcmpl-{completion_id}" if completion_id else None
  14. self.object: str = "chat.completion"
  15. self.created: int = created
  16. self.model: str = None
  17. self.provider: str = None
  18. self.choices = [ChatCompletionChoice(ChatCompletionMessage(content), finish_reason)]
  19. self.usage: dict[str, int] = {
  20. "prompt_tokens": 0, #prompt_tokens,
  21. "completion_tokens": 0, #completion_tokens,
  22. "total_tokens": 0, #prompt_tokens + completion_tokens,
  23. }
  24. def to_json(self):
  25. return {
  26. **self.__dict__,
  27. "choices": [choice.to_json() for choice in self.choices]
  28. }
  29. class ChatCompletionChunk(Model):
  30. def __init__(
  31. self,
  32. content: str,
  33. finish_reason: str,
  34. completion_id: str = None,
  35. created: int = None
  36. ):
  37. self.id: str = f"chatcmpl-{completion_id}" if completion_id else None
  38. self.object: str = "chat.completion.chunk"
  39. self.created: int = created
  40. self.model: str = None
  41. self.provider: str = None
  42. self.choices = [ChatCompletionDeltaChoice(ChatCompletionDelta(content), finish_reason)]
  43. def to_json(self):
  44. return {
  45. **self.__dict__,
  46. "choices": [choice.to_json() for choice in self.choices]
  47. }
  48. class ChatCompletionMessage(Model):
  49. def __init__(self, content: Union[str, None]):
  50. self.role = "assistant"
  51. self.content = content
  52. def to_json(self):
  53. return self.__dict__
  54. class ChatCompletionChoice(Model):
  55. def __init__(self, message: ChatCompletionMessage, finish_reason: str):
  56. self.index = 0
  57. self.message = message
  58. self.finish_reason = finish_reason
  59. def to_json(self):
  60. return {
  61. **self.__dict__,
  62. "message": self.message.to_json()
  63. }
  64. class ChatCompletionDelta(Model):
  65. content: Union[str, None] = None
  66. def __init__(self, content: Union[str, None]):
  67. if content is not None:
  68. self.content = content
  69. def to_json(self):
  70. return self.__dict__
  71. class ChatCompletionDeltaChoice(Model):
  72. def __init__(self, delta: ChatCompletionDelta, finish_reason: Union[str, None]):
  73. self.delta = delta
  74. self.finish_reason = finish_reason
  75. def to_json(self):
  76. return {
  77. **self.__dict__,
  78. "delta": self.delta.to_json()
  79. }
  80. class Image(Model):
  81. url: str
  82. def __init__(self, url: str) -> None:
  83. self.url = url
  84. class ImagesResponse(Model):
  85. data: list[Image]
  86. def __init__(self, data: list) -> None:
  87. self.data = data