tasks.py

from celery import Celery
import google.generativeai as genai
import io
import json
import requests
import os
import logging
from dotenv import load_dotenv
from models import SessionLocal, ApiCall, TaskStatus
from PyPDF2 import PdfReader
from docx import Document
import tempfile

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s',
    force=True
)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# Add a stream handler to ensure logs go to stdout
handler = logging.StreamHandler()
handler.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

load_dotenv()

# Configure Celery with Redis settings
REDIS_HOST = os.getenv('REDIS_HOST', 'redis')
REDIS_PORT = os.getenv('REDIS_PORT', '6379')
celery_app = Celery(
    'gemini_tasks', 
    broker=f'redis://{REDIS_HOST}:{REDIS_PORT}/0',
    backend=f'redis://{REDIS_HOST}:{REDIS_PORT}/0'
)

# Enable remote control and task events to support inspect methods
celery_app.conf.update(
    worker_enable_remote_control=True,
    worker_send_task_events=True,
    worker_pool_restarts=True  # Optional: allows worker pool to restart gracefully
)

# Configure Gemini
genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
DEFAULT_MODEL = os.getenv('GEMINI_MODEL', 'gemini-1.5-flash')
VISION_MODEL = os.getenv('GEMINI_VISION_MODEL', DEFAULT_MODEL)
VISION_PROMPT = os.getenv('GEMINI_VISION_PROMPT', '')

def store_api_call(endpoint, prompt, data, response, tokens_used, file_name=None, file_type=None, status=TaskStatus.PENDING, extracted_text=None, task_id=None):
    db = SessionLocal()
    try:
        api_call = ApiCall(
            endpoint=endpoint,
            prompt=prompt,
            data=data,
            response=response,
            tokens_used=tokens_used,
            file_name=file_name,
            file_type=file_type,
            status=status,
            extracted_text=extracted_text,
            task_id=task_id
        )
        db.add(api_call)
        db.commit()
        return api_call.id
    finally:
        db.close()

def update_api_call(record_id, response=None, status=None, tokens_used=None, extracted_text=None):
    db = SessionLocal()
    try:
        api_call = db.query(ApiCall).filter(ApiCall.id == record_id).first()
        if api_call:
            if response is not None:
                api_call.response = response
            if status is not None:
                api_call.status = status
            if tokens_used is not None:
                api_call.tokens_used = tokens_used
            if extracted_text is not None:
                api_call.extracted_text = extracted_text
            db.commit()
    finally:
        db.close()

def send_callback(callback_url, response_data):
    if callback_url:
        try:
            logger.debug(f"Attempting to send callback to: {callback_url}")
            logger.debug(f"Callback data: {response_data}")
            logger.debug("Making POST request...")
            response = requests.post(callback_url, json=response_data, timeout=10)
            logger.debug(f"Callback response status: {response.status_code}")
            logger.debug(f"Callback response text: {response.text}")
            logger.debug(f"Callback response headers: {response.headers}")
        except Exception as e:
            logger.error(f"Error sending callback to {callback_url}: {str(e)}")
            logger.error(f"Error type: {type(e)}")
            logger.error(f"Error details: {repr(e)}")

def extract_text_from_document(file_contents, file_type):
    """Extract text from various document formats."""
    try:
        if file_type == 'application/pdf':
            # Handle PDF
            pdf_file = io.BytesIO(file_contents)
            pdf_reader = PdfReader(pdf_file)
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text() + "\n"
            return text

        elif file_type in ['application/msword', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document']:
            # Handle DOC/DOCX
            with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as temp_file:
                temp_file.write(file_contents)
                temp_file.flush()
                doc = Document(temp_file.name)
                text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
                os.unlink(temp_file.name)  # Delete the temporary file
                return text

        elif file_type == 'application/rtf':
            # For RTF, convert to plain text (basic approach)
            text = file_contents.decode('utf-8', errors='ignore')
            return text

        else:
            raise ValueError(f"Unsupported file type: {file_type}")

    except Exception as e:
        raise Exception(f"Error extracting text from document: {str(e)}")

@celery_app.task
def process_task(prompt, data=None, callback_url=None, record_id=None):
    try:
        logger.debug(f"Processing task with prompt: {prompt}")
        logger.debug(f"Callback URL: {callback_url}")
        logger.debug(f"Record ID: {record_id}")

        # Update API call with processing status
        if record_id:
            with SessionLocal() as db:
                api_call = db.query(ApiCall).filter(ApiCall.id == record_id).first()
                if api_call:
                    api_call.status = TaskStatus.PROCESSING
                    db.commit()

        # Generate response using model
        model = genai.GenerativeModel(DEFAULT_MODEL)
        response = model.generate_content(json.dumps(prompt))
        response_data = {"response": response.text}
        logger.debug(f"Generated response: {response_data}")

        # Update API call with response
        if record_id:
            with SessionLocal() as db:
                api_call = db.query(ApiCall).filter(ApiCall.id == record_id).first()
                if api_call:
                    api_call.status = TaskStatus.COMPLETED
                    api_call.response = response_data
                    db.commit()

        # Send callback if URL provided
        if callback_url:
            send_callback(callback_url, response_data)

        return response_data
    except Exception as e:
        error_response = {"error": str(e)}
        logger.error(f"Error in process_task: {str(e)}")
        logger.error(f"Error type: {type(e)}")
        logger.error(f"Error details: {repr(e)}")

        # Update API call with error
        if record_id:
            with SessionLocal() as db:
                api_call = db.query(ApiCall).filter(ApiCall.id == record_id).first()
                if api_call:
                    api_call.status = TaskStatus.ERROR
                    api_call.error = str(e)
                    db.commit()

        return error_response

@celery_app.task
def process_document(prompt, file_contents, file_name, file_type, callback_url=None, record_id=None):
    try:
        # Extract text from document
        document_text = extract_text_from_document(file_contents, file_type)
        logger.debug(f"Extracted text length: {len(document_text)} characters")
        
        # Update API call with extracted text
        if record_id:
            update_api_call(
                record_id=record_id,
                extracted_text=document_text
            )
        
        # Use environment vision prompt if available
        final_prompt = VISION_PROMPT if VISION_PROMPT else json.dumps(prompt)
        logger.debug(f"Using prompt: {final_prompt}")
        
        # Generate response using model
        model = genai.GenerativeModel(DEFAULT_MODEL)
        response = model.generate_content([
            f"{final_prompt}\n\nDocument content:\n{document_text}"
        ])
        
        response_data = {"response": response.text}
        logger.debug(f"Generated response: {response_data}")
        
        # Update API call with response
        if record_id:
            update_api_call(
                record_id=record_id,
                response=response_data,
                status=TaskStatus.COMPLETED,
                tokens_used=len(response.text.split())
            )
        
        # Send callback if URL provided
        send_callback(callback_url, response_data)
        
        return response_data
    except Exception as e:
        error_response = {"error": str(e)}
        logger.error(f"Error in process_document: {str(e)}")
        
        # Update API call with error
        if record_id:
            update_api_call(
                record_id=record_id,
                response=error_response,
                status=TaskStatus.FAILED
            )
        
        if callback_url:
            send_callback(callback_url, error_response)
        return error_response