Address comment in openai file

ScottLogic · Dec 19, 2024 · 05695ab · 05695ab
1 parent 103e029
commit 05695ab
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 25 deletions.
diff --git a/backend/src/llm/mistral.py b/backend/src/llm/mistral.py
@@ -42,14 +42,11 @@ async def chat_with_file(
         user_prompt: str,
         files: list[LLMFile],
     ) -> str:
-            try:
-                for file in files:
-                    file = handle_file_upload(files[0])
-                    extracted_content = file["content"]
-                    user_prompt += f"\n\nDocument:\n{extracted_content}"
-                return await self.chat(model, system_prompt, user_prompt)
-            except Exception as file_error:
-                raise HTTPException(
-                    status_code=500,
-                    detail=f"Failed to process files: {str(file_error)}"
-                ) from file_error
+        try:
+            for file in files:
+                file = handle_file_upload(file)
+                extracted_content = file["content"]
+                user_prompt += f"\n\nDocument:\n{extracted_content}"
+            return await self.chat(model, system_prompt, user_prompt)
+        except Exception as file_error:
+            raise HTTPException(status_code=500, detail=f"Failed to process files: {str(file_error)}") from file_error
diff --git a/backend/src/llm/openai.py b/backend/src/llm/openai.py
@@ -1,6 +1,4 @@
 import logging
-from os import PathLike
-from pathlib import Path
 
 from src.utils import Config
 from src.llm import LLM, LLMFile
@@ -88,17 +86,8 @@ async def __upload_files(self, files: list[LLMFile]) -> list[str]:
         file_ids = []
         for file in files:
             logger.info(f"Uploading file '{file.file_name}' to OpenAI")
-            if isinstance(file.file, (PathLike, str)):
-                file_path = Path(file.file)
-                with file_path.open("rb") as f:
-                    file_bytes = f.read()
-            elif isinstance(file.file, bytes):
-                file_bytes = file.file
-            else:
-                logger.error(f"Unsupported file type for '{file.file_name}'")
-                continue
-            file = await client.files.create(file=(file.file_name, file_bytes), purpose="assistants")
-            # file = await client.files.create(file=(file.file_name, file.file), purpose="assistants")
-            file_ids.append(file.id)
+            file = (file.file_name, file.file) if isinstance(file.file, bytes) else file.file
+            response = await client.files.create(file=file, purpose="assistants")
+            file_ids.append(response.id)
 
         return file_ids