LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2026-05-04 23:52:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

(feat:attach) simplify the format for files

Browse Source
This commit is contained in:
ManishMadan2882
2025-04-11 01:52:42 +05:30
parent 292257770c
commit dfe6a8d3e3
1 changed files with 16 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
application/llm/google_ai.py
View File

@@ -3,6 +3,7 @@ from google.genai import types
import os import os
import logging import logging
import mimetypes import mimetypes
import json
from application.llm.base import BaseLLM from application.llm.base import BaseLLM
@@ -41,11 +42,11 @@ class GoogleLLM(BaseLLM):
Returns: Returns:
list: Messages formatted with file references for Google AI API. list: Messages formatted with file references for Google AI API.
""" """
if not attachments: if not attachments:
return messages return messages
prepared_messages = messages.copy() prepared_messages = messages.copy()
logging.info(f"GoogleLLM: Initial messages before attachment processing: {json.dumps(prepared_messages, indent=2)}")
# Find the user message to attach files to the last one # Find the user message to attach files to the last one
user_message_index = None user_message_index = None
@@ -54,7 +55,6 @@ class GoogleLLM(BaseLLM):
user_message_index = i user_message_index = i
break break
if user_message_index is None: if user_message_index is None:
user_message = {"role": "user", "content": []} user_message = {"role": "user", "content": []}
prepared_messages.append(user_message) prepared_messages.append(user_message)
@@ -68,7 +68,7 @@ class GoogleLLM(BaseLLM):
elif not isinstance(prepared_messages[user_message_index].get("content"), list): elif not isinstance(prepared_messages[user_message_index].get("content"), list):
prepared_messages[user_message_index]["content"] = [] prepared_messages[user_message_index]["content"] = []
file_uris = [] files = []
for attachment in attachments: for attachment in attachments:
mime_type = attachment.get('mime_type') mime_type = attachment.get('mime_type')
if not mime_type: if not mime_type:
@@ -80,7 +80,7 @@ class GoogleLLM(BaseLLM):
try: try:
file_uri = self._upload_file_to_google(attachment) file_uri = self._upload_file_to_google(attachment)
logging.info(f"GoogleLLM: Successfully uploaded file, got URI: {file_uri}") logging.info(f"GoogleLLM: Successfully uploaded file, got URI: {file_uri}")
file_uris.append((file_uri, mime_type)) files.append({"file_uri": file_uri, "mime_type": mime_type})
except Exception as e: except Exception as e:
logging.error(f"GoogleLLM: Error uploading file: {e}") logging.error(f"GoogleLLM: Error uploading file: {e}")
if 'content' in attachment: if 'content' in attachment:
@@ -89,13 +89,13 @@ class GoogleLLM(BaseLLM):
"text": f"[File could not be processed: {attachment.get('path', 'unknown')}]" "text": f"[File could not be processed: {attachment.get('path', 'unknown')}]"
}) })
if file_uris: if files:
logging.info(f"GoogleLLM: Adding {len(file_uris)} file URIs to message") logging.info(f"GoogleLLM: Adding {len(files)} files to message")
prepared_messages[user_message_index]["content"].append({ prepared_messages[user_message_index]["content"].append({
"type": "file_uris", "files": files
"file_uris": file_uris
}) })
logging.info(f"GoogleLLM: Final prepared messages: {json.dumps(prepared_messages, indent=2)}")
return prepared_messages return prepared_messages
def _upload_file_to_google(self, attachment): def _upload_file_to_google(self, attachment):
@@ -149,6 +149,7 @@ class GoogleLLM(BaseLLM):
raise raise
def _clean_messages_google(self, messages): def _clean_messages_google(self, messages):
logging.info(f"GoogleLLM: Starting message cleaning. Input messages: {json.dumps(messages, indent=2)}")
cleaned_messages = [] cleaned_messages = []
for message in messages: for message in messages:
role = message.get("role") role = message.get("role")
@@ -179,14 +180,14 @@ class GoogleLLM(BaseLLM):
response=item["function_response"]["response"], response=item["function_response"]["response"],
) )
) )
elif "type" in item and item["type"] == "file_uris": elif "files" in item:
for file_uri, mime_type in item["file_uris"]: for file_data in item["files"]:
parts.append( parts.append(
types.Part.from_uri( types.Part.from_uri(
file_uri=file_uri, file_uri=file_data["file_uri"],
mime_type=mime_type mime_type=file_data["mime_type"]
)
) )
)
else: else:
raise ValueError( raise ValueError(
f"Unexpected content dictionary format:{item}" f"Unexpected content dictionary format:{item}"