Update app.py

app.py

@@ -16,7 +16,10 @@ from typing import List, Optional, Dict
 from google.oauth2 import service_account
 from googleapiclient.discovery import build
 
+from io import BytesIO
 import google.generativeai as genai
+import os
+import requests
 
 # start application
 app = FastAPI()
@@ -129,71 +132,84 @@ def convert_dicts_to_list(dicts: List[Dict], headers: List[str], add_link: bool
         rows.append(row)
     # print(rows)
     return rows
-    
-async def request_gpt4o_completion(image_id: str, user_credit: int):
-    openai_url = "https://api.openai.com/v1/chat/completions"
-    headers = {"Authorization": f"Bearer {OPENAI_KEY}", "Content-Type": "application/json"}
-    image_url = f"https://drive.usercontent.google.com/u/0/uc?id={image_id}&export=download"
-    body = {
-    "model": "gpt-4o-mini",
-    "response_format": { "type": "json_object" },
-    "messages": [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "text",
-                    "text": f'''
-아래의 순서대로 실행해주세요:
-1. {image_url} 이 링크에서 파일 다운로드가 완료될 때까지 기다리세요.
-2. 이미지의 모든 영수증과 명함을 식별하고, 각각의 정보를 추출하세요.
-3. 작성할 때 영수증과 명함 개수의 합이 {user_credit}을 넘지 않도록 하세요.
-4. 추출한 정보를 {receipt_sheet_headers}란이 있는 영수증 시트 또는 {business_card_sheet_headers}란이 있는 명함 시트에 넣을 수 있도록 JSON으로 정리하세요.
-json.receipts.forEach() 또는 json.busi_cards.forEach()
-(item) => insert row where header field is item.[fieldName]
-'''
-                },
-                {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": image_url,
-                        "detail": "high"
-                    }
-                }
-            ]
-        }
-    ],
-    "max_tokens": 4096
-}
 
-    logging.info(f"Request URL: {openai_url}")
-    logging.info(f"Request Headers: {headers}")
-    logging.info(f"Request Body: {json.dumps(body, indent=2)}")
 
+async def request_gpt4o_completion(image_id: str, user_credit: int):
     try:
-        async with httpx.AsyncClient(follow_redirects=True, timeout=300.0) as client:
-            response = await client.post(openai_url, headers=headers, json=body)
-            response.raise_for_status()  # Raise an error for bad status codes
-    except httpx.RequestError as exc:
-        logging.error(f"An error occurred while requesting {exc.request.url!r}.")
-        logging.error(f"Exception: {str(exc)}")
-        return None
-    except httpx.HTTPStatusError as exc:
-        logging.error(f"HTTP error occurred: {str(exc)}")
-        logging.error(f"Response: {exc.response.text}")
-        return None
+        # Download file from Google Drive
+        file_content, mime_type = download_file_from_drive(image_id)
+        
+        temp_file_path = "/tmp/downloaded_file.pdf"
+        with open(temp_file_path, "wb") as f:
+            f.write(file_content.getbuffer())
+
+        # Log the temporary file path
+        logging.info(f"File saved to {temp_file_path}")
+
+        genai_file = genai.upload_file(path=temp_file_path, display_name="yungsoogi")
+
+        # Clean up: remove the temporary file after processing
+        os.remove(temp_file_path)
 
-    logging.info(f"Response Status Code: {response.status_code}")
-    logging.info(f"Response Body: {response.text}")
         
-    completion = response.json()
-    if 'choices' in completion and len(completion['choices']) > 0:
-        message_content = completion['choices'][0]['message']['content']
-        print(message_content)
-        return message_content
-    else:
-        print("No response received")
-        return None
+        prompt = f'''
+                    아래의 순서대로 실행해주세요:
+                    1. {image_url} 이 링크에서 파일 다운로드가 완료될 때까지 기다리세요.
+                    2. 이미지의 모든 영수증과 명함을 식별하고, 각각의 정보를 추출하세요.
+                    3. 작성할 때 영수증과 명함 개수의 합이 {user_credit}을 넘지 않도록 하세요.
+                    4. 추출한 정보를 {receipt_sheet_headers}란이 있는 영수증 시트 또는 {business_card_sheet_headers}란이 있는 명함 시트에 넣을 수 있도록 JSON으로 정리하세요.
+                    json.receipts.forEach() 또는 json.busi_cards.forEach()
+                    (item) => insert row where header field is item.[fieldName]
+                '''
+        
+        
+        # Generate content using the model
+        ai_response = model.generate_content(
+            [prompt, genai_file]
+        )
+
+    genai.delete_file(response.name)
+
+    logging.info(f"Response Body: {ai_response}")
+
+    return None
+        
+    # if 'choices' in completion and len(completion['choices']) > 0:
+    #     message_content = completion['choices'][0]['message']['content']
+    #     print(message_content)
+    #     return message_content
+    # else:
+    #     print("No response received")
+    #     return None
+
+def download_file_from_drive(file_id: str) -> (BytesIO, str):
+    """
+    Downloads a file from Google Drive using the file ID
+    and returns a BytesIO stream and the MIME type.
+    """
+    try:
+        # URL for downloading the file from Google Drive
+        url = f"https://drive.google.com/uc?export=download&id={file_id}"
+
+        # Send a GET request to the URL
+        response = requests.get(url, stream=True)
+
+        # Check if the request was successful
+        if response.status_code == 200:
+            # Create a BytesIO object from the response content
+            file_io = BytesIO(response.content)
+
+            # Get the MIME type of the file from the response headers
+            mime_type = response.headers.get('Content-Type', 'application/octet-stream')
+
+            return file_io, mime_type
+        else:
+            raise HTTPException(status_code=500, detail=f"Failed to download file, status code: {response.status_code}")
+
+    except Exception as e:
+        logging.error(f"Error downloading file from Drive: {e}")
+        raise HTTPException(status_code=500, detail="Failed to download file from Drive")
+
     
 async def move_file_to_folder(file_id, current_parents, new_parents):
     try: