app.py

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)