🗜️ squash schemas

shroominic · Dec 12, 2023 · ac9f2ea · ac9f2ea
1 parent f88be8c
commit ac9f2ea
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 69 deletions.
diff --git a/src/codeinterpreterapi/schema/file.py → src/codeinterpreterapi/schema.py b/src/codeinterpreterapi/schema/file.py → src/codeinterpreterapi/schema.py
@@ -1,7 +1,9 @@
 import asyncio
 
 from typing import Any
+from langchain.schema import AIMessage, HumanMessage
 from langchain.pydantic_v1 import BaseModel
+from codeboxapi.schema import CodeBoxStatus
 
 
 class File(BaseModel):
@@ -87,3 +89,54 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         return f"File(name={self.name})"
+
+
+class CodeInput(BaseModel):
+    code: str
+
+
+class FileInput(BaseModel):
+    filename: str
+
+
+class UserRequest(HumanMessage):
+    files: list[File] = []
+
+    def __str__(self) -> str:
+        return str(self.content)
+
+    def __repr__(self) -> str:
+        return f"UserRequest(content={self.content}, files={self.files})"
+
+
+class CodeInterpreterResponse(AIMessage):
+    """
+    Response from the code interpreter agent.
+
+    files: list of files to be sent to the user (File )
+    code_log: list[tuple[str, str]] = []
+    """
+
+    files: list[File] = []
+    code_log: list[tuple[str, str]] = []
+
+    def show(self) -> None:
+        print("AI: ", self.content)
+        for file in self.files:
+            print("File: ", file.name)
+            file.show_image()
+
+    def __str__(self) -> str:
+        return str(self.content)
+
+    def __repr__(self) -> str:
+        return f"CodeInterpreterResponse(content={self.content}, files={self.files})"
+
+
+class SessionStatus(CodeBoxStatus):
+    @classmethod
+    def from_codebox_status(cls, cbs: CodeBoxStatus) -> "SessionStatus":
+        return cls(status=cbs.status)
+
+    def __repr__(self) -> str:
+        return f"<SessionStatus status={self.status}>"
diff --git a/src/codeinterpreterapi/schema/__init__.py b/src/codeinterpreterapi/schema/__init__.py
diff --git a/src/codeinterpreterapi/schema/input.py b/src/codeinterpreterapi/schema/input.py
diff --git a/src/codeinterpreterapi/schema/response.py b/src/codeinterpreterapi/schema/response.py
diff --git a/src/codeinterpreterapi/schema/status.py b/src/codeinterpreterapi/schema/status.py