Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import requests | |
| import tempfile | |
| from google.oauth2 import service_account | |
| from googleapiclient.discovery import build | |
| from googleapiclient.http import MediaIoBaseDownload | |
| import openai | |
| from dotenv import load_dotenv, dotenv_values | |
| import io | |
| from markitdown import MarkItDown | |
| from openai import OpenAI | |
| openai.api_key = os.getenv('OPENAI_API_KEY') | |
| openai = OpenAI(api_key = openai.api_key) | |
| class GPTDriveIntegration: | |
| def __init__(self): | |
| # Build credentials info from individual environment variables | |
| credentials_info = { | |
| "type": "service_account", | |
| "project_id": os.getenv('GOOGLE_PROJECT_ID'), | |
| "private_key_id": os.getenv('GOOGLE_PRIVATE_KEY_ID'), | |
| "private_key": os.getenv('GOOGLE_PRIVATE_KEY').replace('\\n', '\n'), # Fix line breaks | |
| "client_email": os.getenv('GOOGLE_CLIENT_EMAIL'), | |
| "client_id": os.getenv('GOOGLE_CLIENT_ID'), | |
| "auth_uri": "https://accounts.google.com/o/oauth2/auth", | |
| "token_uri": "https://oauth2.googleapis.com/token", | |
| "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", | |
| "client_x509_cert_url": os.getenv('GOOGLE_CLIENT_CERT_URL'), | |
| "universe_domain": "googleapis.com" | |
| } | |
| # Check if all required fields are present | |
| required_fields = ['project_id', 'private_key', 'client_email'] | |
| missing_fields = [field for field in required_fields if not credentials_info[field]] | |
| if missing_fields: | |
| raise ValueError(f"Missing required environment variables: {missing_fields}") | |
| # Initialize Google Drive API | |
| self.credentials = service_account.Credentials.from_service_account_info( | |
| credentials_info, | |
| scopes=['https://www.googleapis.com/auth/drive.readonly'] | |
| ) | |
| self.drive_service = build('drive', 'v3', credentials=self.credentials) | |
| # Initialize MarkItDown | |
| self.md = MarkItDown() | |
| # Initialize OpenAI | |
| openai.api_key = os.getenv('OPENAI_API_KEY') | |
| def search_files(self, query, file_types=None): | |
| """Search for files in Google Drive""" | |
| search_query = f"name contains '{query}'" | |
| if file_types: | |
| type_queries = [] | |
| for file_type in file_types: | |
| ext = file_type.lower().lstrip('.') | |
| if ext == 'pdf': | |
| type_queries.append("mimeType='application/pdf'") | |
| elif ext in ['doc', 'docx']: | |
| type_queries.append("mimeType contains 'document'") | |
| elif ext in ['xls', 'xlsx']: | |
| type_queries.append("mimeType contains 'spreadsheet'") | |
| elif ext in ['ppt', 'pptx']: | |
| type_queries.append("mimeType contains 'presentation'") | |
| elif ext in ['txt', 'md', 'markdown']: | |
| type_queries.append("mimeType='text/plain'") | |
| if type_queries: | |
| search_query += f" and ({' or '.join(type_queries)})" | |
| results = self.drive_service.files().list( | |
| q=search_query, | |
| fields="files(id, name, mimeType, size)" | |
| ).execute() | |
| return results.get('files', []) | |
| def get_file_content(self, file_id, file_name, mime_type): | |
| """Download and extract content from file using MarkItDown""" | |
| try: | |
| # Handle Google Workspace files - export to appropriate format for MarkItDown | |
| if 'document' in mime_type: | |
| # Export Google Docs as DOCX for better formatting preservation | |
| request = self.drive_service.files().export_media( | |
| fileId=file_id, | |
| mimeType='application/vnd.openxmlformats-officedocument.wordprocessingml.document' | |
| ) | |
| file_extension = 'docx' | |
| elif 'spreadsheet' in mime_type: | |
| # Export Google Sheets as XLSX | |
| request = self.drive_service.files().export_media( | |
| fileId=file_id, | |
| mimeType='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet' | |
| ) | |
| file_extension = 'xlsx' | |
| elif 'presentation' in mime_type: | |
| # Export Google Slides as PPTX | |
| request = self.drive_service.files().export_media( | |
| fileId=file_id, | |
| mimeType='application/vnd.openxmlformats-officedocument.presentationml.presentation' | |
| ) | |
| file_extension = 'pptx' | |
| else: | |
| # For regular files, download as-is | |
| request = self.drive_service.files().get_media(fileId=file_id) | |
| file_extension = self._get_extension_from_name_or_mime(file_name, mime_type) | |
| # Download file content | |
| file_content = io.BytesIO() | |
| downloader = MediaIoBaseDownload(file_content, request) | |
| done = False | |
| while done is False: | |
| status, done = downloader.next_chunk() | |
| # Reset stream position | |
| file_content.seek(0) | |
| # Use MarkItDown to convert to markdown | |
| result = self.md.convert_stream(file_content, file_extension=file_extension) | |
| return result.text_content | |
| except Exception as e: | |
| return f"Error processing file with MarkItDown: {str(e)}" | |
| def _get_extension_from_name_or_mime(self, file_name, mime_type): | |
| """Helper to determine file extension for MarkItDown""" | |
| # First try to get extension from filename | |
| if '.' in file_name: | |
| return file_name.split('.')[-1].lower() | |
| # Fallback to mime type mapping | |
| mime_to_ext = { | |
| 'application/pdf': 'pdf', | |
| 'application/vnd.openxmlformats-officedocument.wordprocessingml.document': 'docx', | |
| 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': 'xlsx', | |
| 'application/vnd.openxmlformats-officedocument.presentationml.presentation': 'pptx', | |
| 'application/msword': 'doc', | |
| 'application/vnd.ms-excel': 'xls', | |
| 'application/vnd.ms-powerpoint': 'ppt', | |
| 'text/plain': 'txt', | |
| 'text/markdown': 'md', | |
| 'text/html': 'html', | |
| 'application/json': 'json', | |
| 'text/csv': 'csv' | |
| } | |
| return mime_to_ext.get(mime_type, 'txt') | |
| def query_gpt_with_context(self, user_query, file_contents): | |
| """Send query to GPT with file context""" | |
| context = "\n\n".join([ | |
| f"File: {content['name']}\nContent: {content['text'][:3000]}..." | |
| for content in file_contents | |
| ]) | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": """ | |
| You are an AI assistant that can analyze documents from Google Drive. | |
| Use the provided file contents to answer user questions. | |
| Answer directly and add additional suggestions on how to answer questions in the exam | |
| Always end with 'Is there anything I can help you with?' | |
| Your name is Study buddy, happy to help students study more effectively | |
| """ | |
| }, | |
| { | |
| "role": "user", | |
| "content": f"Context from Google Drive files:\n{context}\n\nUser Question: {user_query}" | |
| } | |
| ] | |
| response = openai.chat.completions.create( | |
| model="gpt-4o-mini", | |
| messages=messages, | |
| max_tokens=1000 | |
| ) | |
| return response.choices[0].message.content | |
| def process_query(self, user_query, search_terms=None): | |
| """Main function to process user queries""" | |
| # Extract search terms from query if not provided | |
| if not search_terms: | |
| search_terms = user_query.split()[:3] # Simple extraction | |
| # Search for relevant files | |
| files = [] | |
| for term in search_terms: | |
| files.extend(self.search_files(term)) | |
| # Remove duplicates | |
| unique_files = {f['id']: f for f in files}.values() | |
| # Get content from top 3 most relevant files | |
| file_contents = [] | |
| for file in list(unique_files)[:3]: | |
| content = self.get_file_content(file['id'], file['name'], file['mimeType']) | |
| file_contents.append({ | |
| 'name': file['name'], | |
| 'text': content | |
| }) | |
| # Query GPT with context | |
| if file_contents: | |
| response = self.query_gpt_with_context(user_query, file_contents) | |
| return { | |
| 'answer': response, | |
| 'sources': [f['name'] for f in file_contents] | |
| } | |
| else: | |
| return { | |
| 'answer': "No relevant files found in your Google Drive.", | |
| 'sources': [] | |
| } | |
| gpt_drive = GPTDriveIntegration() | |
| def process_user_query(query, search_terms_input): | |
| """Process user query and return formatted response""" | |
| if not query.strip(): | |
| return "Please enter a question.", "" | |
| # Parse search terms if provided | |
| search_terms = None | |
| # if search_terms_input.strip(): | |
| # search_terms = [term.strip() for term in search_terms_input.split(',')] | |
| # Process the query | |
| result = gpt_drive.process_query(query, search_terms) | |
| # Format the response | |
| answer = result['answer'] | |
| sources = result['sources'] | |
| sources_text = "" | |
| if sources: | |
| sources_text = "**Sources used:**\n" + "\n".join([f"• {source}" for source in sources]) | |
| return answer, sources_text | |
| def check_setup(): | |
| """Check if the APIs are properly configured""" | |
| status_messages = [] | |
| # Check Google Drive API | |
| if hasattr(gpt_drive, 'drive_initialized') and gpt_drive.drive_initialized: | |
| status_messages.append("✅ Google Drive API: Connected") | |
| else: | |
| status_messages.append(f"❌ Google Drive API: {getattr(gpt_drive, 'drive_error', 'Not configured')}") | |
| # Check OpenAI API | |
| if hasattr(gpt_drive, 'openai_initialized') and gpt_drive.openai_initialized: | |
| status_messages.append("✅ OpenAI API: Connected") | |
| else: | |
| status_messages.append(f"❌ OpenAI API: {getattr(gpt_drive, 'openai_error', 'Not configured')}") | |
| return "\n".join(status_messages) | |
| # Create Gradio interface | |
| import gradio as gr | |
| with gr.Blocks(title="Study Buddy", theme=gr.themes.Soft()) as app: | |
| gr.Markdown("# Anatomy Study Buddy ") | |
| gr.Markdown("Study more effectively with study Buddy!") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| # Main query interface | |
| with gr.Group(): | |
| gr.Markdown("### Ask a Question") | |
| query_input = gr.Textbox( | |
| label="Your Question", | |
| placeholder="Ask me any question about your anatomy books?", | |
| lines=3 | |
| ) | |
| search_terms_input = gr.Textbox( | |
| label="Search Terms", | |
| placeholder="Enter comma-separated terms to search for specific files", | |
| lines=1 | |
| ) | |
| submit_btn = gr.Button("Search & Ask", variant="primary", size="lg") | |
| # Results section | |
| with gr.Group(): | |
| gr.Markdown("### Answer") | |
| answer_output = gr.Textbox( | |
| label="AI Response", | |
| lines=10, | |
| interactive=False | |
| ) | |
| sources_output = gr.Textbox( | |
| label="Sources", | |
| lines=3, | |
| interactive=False | |
| ) | |
| # Event handlers | |
| submit_btn.click( | |
| fn=process_user_query, | |
| inputs=[query_input, search_terms_input], | |
| outputs=[answer_output, sources_output] | |
| ) | |
| # Example queries | |
| with gr.Row(): | |
| gr.Examples( | |
| examples=[ | |
| ["What is morbid Anatomy?", "morbid, Anatomy"], | |
| ["The transmission of nerves from one neuron to another is as a result of what?", "neuron, nerves, Dr Clement"], | |
| ["Explain what the external ear contains of?", "Ear Anatomy, Ear"], | |
| ["What are the types of massage?", "massage Lecture, nerves"], | |
| ["What is trauma?", "Trauma, physical trauma and sex Offenders"], | |
| ["what is Upper limb prosthetics?", "Upper limb prosthetics"], | |
| ], | |
| inputs=[query_input, search_terms_input], | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| app.launch( | |
| share=True,debug =True) |