finalize schemas

codeinterpreterapi/schemas.py

@@ -1,13 +1,26 @@
-from promptkit import AssistantResponse
-from promptkit.schema import UserMessage
+from langchain.schema import HumanMessage, AIMessage  # type: ignore
+from pydantic import BaseModel
 
 
-class UserRequest(UserMessage):
+class File(BaseModel):
+    name: str
+    content: bytes
+
+
+class UserRequest(HumanMessage):
     text: str = ""
-    files: list[dict[str, bytes]] = []
+    files: list[File] = []
 
 
-class CodeInterpreterResponse(AssistantResponse):
-    files: list[dict[str, bytes]] = []
+class CodeInterpreterResponse(AIMessage):
+    files: list[File] = []
     final_code: str = ""
     final_output: str = ""
+
+
+class CodeInput(BaseModel): 
+    code: str
+    
+
+class FileInput(BaseModel): 
+    filename: str