61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import uuid
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.tasks.celery_app import celery_app
|
|
from app.tasks.transcription import SyncSession
|
|
from app.models.protocol import Protocol, ProtocolStatus
|
|
from app.models.transcription import Transcription
|
|
from app.services.protocol_generator import (
|
|
generate_protocol_text,
|
|
get_default_template,
|
|
get_template_by_id,
|
|
)
|
|
|
|
|
|
@celery_app.task(bind=True, name="app.tasks.protocol.generate_protocol")
|
|
def generate_protocol(self, transcription_id: str, template_id: str | None = None):
|
|
trans_uuid = uuid.UUID(transcription_id)
|
|
|
|
with SyncSession() as db:
|
|
transcription = db.execute(
|
|
select(Transcription).where(Transcription.id == trans_uuid)
|
|
).scalar_one_or_none()
|
|
|
|
if not transcription:
|
|
return {"error": "Transcription not found"}
|
|
|
|
if template_id:
|
|
template = get_template_by_id(db, uuid.UUID(template_id))
|
|
else:
|
|
template = get_default_template(db)
|
|
|
|
if not template:
|
|
return {"error": "No template found"}
|
|
|
|
protocol = Protocol(
|
|
transcription_id=trans_uuid,
|
|
template_id=template.id,
|
|
llm_model="",
|
|
status=ProtocolStatus.generating,
|
|
)
|
|
db.add(protocol)
|
|
db.commit()
|
|
db.refresh(protocol)
|
|
|
|
try:
|
|
result = generate_protocol_text(transcription.text, template)
|
|
|
|
protocol.content = result["content"]
|
|
protocol.raw_text = result["raw_text"]
|
|
protocol.llm_model = result["model"]
|
|
protocol.status = ProtocolStatus.done
|
|
db.commit()
|
|
|
|
return {"protocol_id": str(protocol.id), "status": "done"}
|
|
|
|
except Exception as e:
|
|
protocol.status = ProtocolStatus.error
|
|
db.commit()
|
|
raise self.retry(exc=e, max_retries=2, countdown=60)
|