Spaces:
Sleeping
Sleeping
File size: 2,500 Bytes
37eaffd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from __future__ import annotations
import json
from pathlib import Path
from ..inference.endpoint import EndpointCapability
from ..models.schema import TEISchema
from ..models.spans import SpanDescriptor
try:
from jinja2 import Environment, FileSystemLoader
_HAS_JINJA = True
except ImportError:
_HAS_JINJA = False
_TEMPLATE_DIR = Path(__file__).parent / "templates"
def _get_env() -> "Environment":
if not _HAS_JINJA:
raise ImportError(
"jinja2 is required for prompt building. Install it with: pip install jinja2"
)
env = Environment(loader=FileSystemLoader(str(_TEMPLATE_DIR)), keep_trailing_newline=True)
env.filters["tojson"] = lambda x, **kw: json.dumps(x, ensure_ascii=False, **kw)
return env
def build_prompt(
source_text: str,
schema: TEISchema,
capability: EndpointCapability,
candidates: list[SpanDescriptor] | None = None,
) -> str:
"""
Build an LLM prompt for the given endpoint capability.
Raises ValueError for EXTRACTION endpoints (they don't use text prompts).
"""
if capability == EndpointCapability.EXTRACTION:
raise ValueError(
"EXTRACTION endpoints use their own native format; no text prompt needed."
)
env = _get_env()
template_name = (
"text_gen.jinja2"
if capability == EndpointCapability.TEXT_GENERATION
else "json_enforced.jinja2"
)
template = env.get_template(template_name)
candidate_dicts: list[dict] | None = None
if candidates:
candidate_dicts = [
{
"element": c.element,
"text": c.text,
"context": c.context,
"attrs": c.attrs,
**({"confidence": c.confidence} if c.confidence is not None else {}),
}
for c in candidates
]
return template.render(
schema=schema,
source_text=source_text,
candidates=candidate_dicts,
)
def make_correction_prompt(original_response: str, error_message: str) -> str:
"""Build a self-correction retry prompt that includes the bad response and the error."""
return (
"Your previous response could not be parsed as JSON.\n"
f"Error: {error_message}\n\n"
f"Your previous response was:\n{original_response}\n\n"
"Please fix the JSON and return only a valid JSON array of span objects. "
"Do not include any markdown formatting or explanation."
)
|