Skip to content

Commit

Permalink
Improved file classification for tool resources
Browse files Browse the repository at this point in the history
  • Loading branch information
VRSEN committed Aug 23, 2024
1 parent d9c3245 commit 53a0e62
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 13 deletions.
4 changes: 2 additions & 2 deletions agency_swarm/agency/agency.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from agency_swarm.tools import BaseTool, CodeInterpreter, FileSearch
from agency_swarm.user import User
from agency_swarm.util.errors import RefusalError
from agency_swarm.util.files import determine_file_type
from agency_swarm.util.files import determine_file_type, get_tools
from agency_swarm.util.shared_state import SharedState
from agency_swarm.util.streaming import AgencyEventHandler

Expand Down Expand Up @@ -349,7 +349,7 @@ def handle_file_upload(file_list):
else:
attachments.append({
"file_id": file.id,
"tools": tools
"tools": get_tools(file.filename)
})

message_file_names.append(file.filename)
Expand Down
41 changes: 30 additions & 11 deletions agency_swarm/util/files.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,38 @@
import mimetypes

code_interpreter_types = [
"application/csv", "image/jpeg", "image/gif", "image/png",
"application/x-tar", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"application/xml", "text/xml", "application/zip"
]

dual_types = [
"text/x-c", "text/x-csharp", "text/x-c++", "application/msword",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"text/html", "text/x-java", "application/json", "text/markdown",
"application/pdf", "text/x-php",
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
"text/x-python", "text/x-script.python", "text/x-ruby", "text/x-tex",
"text/plain", "text/css", "text/javascript", "application/x-sh",
"application/typescript"
]

def determine_file_type(file_path):
mime_type, _ = mimetypes.guess_type(file_path)
if mime_type:
if mime_type in [
'application/json', 'text/csv', 'application/xml',
'application/vnd.ms-excel', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'application/zip'
]:
if mime_type in code_interpreter_types:
return "assistants.code_interpreter"
elif mime_type in [
'text/plain', 'text/markdown', 'application/pdf',
'application/msword', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
]:
return "assistants.file_search"
elif mime_type.startswith('image/'):
return "vision"
return "assistants.file_search"
elif mime_type in dual_types:
return "assistants.file_search"
raise ValueError(f"Unsupported file type: {mime_type}")

def get_tools(file_path):
mime_type, _ = mimetypes.guess_type(file_path)
if mime_type in code_interpreter_types:
return [{"type": "code_interpreter"}]
elif mime_type in dual_types:
return [{"type": "code_interpreter"}, {"type": "retrieval"}]
else:
raise ValueError(f"Unsupported file type: {mime_type}")

0 comments on commit 53a0e62

Please sign in to comment.