Document Embeddings Pipeline

from hyrex import HyrexRegistry
import psycopg2
import openai
import os

hy = HyrexRegistry()

@hy.task
def update_db_record(doc_id: str, embedding: list[float]):
    """Store document embedding in PostgreSQL"""
    with psycopg2.connect(os.environ.get("DB_CONN_STRING")) as conn:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE documents SET embedding = %s WHERE doc_id = %s",
            (embedding, doc_id)
        )

@hy.task
def process_document(doc_id: str):
    """Generate embedding for a single document"""
    # Fetch document content from Google Drive
    file_content = gdrive_sdk.get_document(doc_id)
    
    # Generate embedding using OpenAI
    response = openai.embeddings.create(
        model="text-embedding-3-small",
        input=file_content
    )
    embedding = response.data[0].embedding
    
    # Store embedding in database
    update_db_record.send(doc_id, embedding)

@hy.task
def sync_google_drive_documents():
    """Process all documents in Google Drive"""
    all_doc_ids = gdrive_sdk.list_document_ids()
    for doc_id in all_doc_ids:
        process_document.send(doc_id)

from fastapi import FastAPI
from pydantic import BaseModel
from .tasks import process_document, sync_google_drive_documents

app = FastAPI()

class DocumentRequest(BaseModel):
    doc_id: str
    
class BulkSyncRequest(BaseModel):
    folder_id: str = None

@app.post("/sync/document")
async def sync_document(request: DocumentRequest):
    """Process a single document"""
    task = process_document.send(request.doc_id)
    return {"message": "Document sync started", "task_id": task.task_id}

@app.post("/sync/all-documents") 
async def sync_all_documents(request: BulkSyncRequest = None):
    """Process all Google Drive documents"""
    task = sync_google_drive_documents.send()
    return {"message": "Bulk sync started", "task_id": task.task_id}