52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
import json
|
|
import uuid
|
|
|
|
from litellm import completion
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.config import settings
|
|
from app.models.prompt_template import PromptTemplate
|
|
|
|
|
|
def generate_protocol_text(transcription_text: str, template: PromptTemplate) -> dict:
|
|
user_message = template.user_prompt.replace("{transcription}", transcription_text)
|
|
|
|
response = completion(
|
|
model=settings.litellm_model,
|
|
api_key=settings.litellm_api_key,
|
|
api_base=settings.litellm_api_base or None,
|
|
messages=[
|
|
{"role": "system", "content": template.system_prompt},
|
|
{"role": "user", "content": user_message},
|
|
],
|
|
temperature=0.3,
|
|
response_format={"type": "json_object"},
|
|
)
|
|
|
|
raw_text = response.choices[0].message.content
|
|
try:
|
|
content = json.loads(raw_text)
|
|
except json.JSONDecodeError:
|
|
content = {"raw": raw_text}
|
|
|
|
return {
|
|
"content": content,
|
|
"raw_text": raw_text,
|
|
"model": response.model,
|
|
}
|
|
|
|
|
|
def get_default_template(db: Session) -> PromptTemplate | None:
|
|
result = db.execute(
|
|
select(PromptTemplate).where(PromptTemplate.is_default == True)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
def get_template_by_id(db: Session, template_id: uuid.UUID) -> PromptTemplate | None:
|
|
result = db.execute(
|
|
select(PromptTemplate).where(PromptTemplate.id == template_id)
|
|
)
|
|
return result.scalar_one_or_none()
|