Structured Output with AzureOpenAI

[Wolftess]

Hi everyone,
first of all, thanks a lot for your great work and for providing this library. I am fascinated by its few-shot capabilities.
Now to my question: I am using a custom provider like this for using AzureOpenAI models:

@lx.providers.registry.register("azure-openai")
class AzureOpenAILanguageModel(lx.core.base_model.BaseLanguageModel):
    # Acquire AAD bearer token for Azure OpenAI
    credential = DefaultAzureCredential()
    token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
    # No direct use, we pass the provider to the SDK
    azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
    if not azure_endpoint:
        raise RuntimeError("AZURE_OPENAI_ENDPOINT environment variable is not set.")
    # Defaults from llm config (can be overridden via language_model_params)

    def __init__(self, azure_endpoint=azure_endpoint, token_provider=token_provider, **kwargs):
        # Extract constraint before passing to parent
        constraint = kwargs.pop("constraint", None)
        
        # Call parent to initialize _schema, _constraint, etc.
        super().__init__(constraint=constraint, **kwargs)
        
        # Azure-specific configuration
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")
        self.max_workers = kwargs.get("max_workers")
        self.api_version = kwargs.get("api_version")
        self.response_format = kwargs.get("response_format")
        
        # Azure deployment/model name
        self.deployment = kwargs.get("deployment_name")

        
        # Azure client
        self.client = AzureOpenAI(
            azure_ad_token_provider=token_provider,
            azure_endpoint=azure_endpoint,
            api_version=self.api_version,
        )

    def infer(self, batch_prompts, **kwargs):
        for prompt in batch_prompts:
            resp = self.client.chat.completions.parse(
                model=self.deployment,
                messages=[{"role": "user", "content": prompt}],
                response_format=self.response_format,  # Pydantic schema
            )
            parsed = resp.choices[0].message.parsed  # a Pydantic model
            output_str = parsed.model_dump_json()    # <- string for LangExtract

            #Wrap the JSON string in the markers LangExtract expects
            #output_with_markers = f"```json\n{output_str}\n```"

            if not output_str.strip() or output_str == "{}":
                # If the output is empty, yield an empty response that LangExtract can handle
                yield [lx.core.types.ScoredOutput(score=1.0, output="")]
            else:
                # Yield the formatted string with markers
                yield [lx.core.types.ScoredOutput(score=1.0, output=output_str)]

This works. However, I am using a special output_format which looks like this:

class BPMNAnnotation(BaseModel):
extraction_class: str = Field(..., description="Entity/type label")
extraction_text: str = Field(..., description="Exact text from source")
attributes: Optional[Dict[str, Union[str, int, float, bool, None]]] = Field(
None, description="Key-value attributes"
)

class BPMNExtractions(BaseModel):
extractions: List[BPMNAnnotation]

I am calling it like this:
result = lx.extract(
text_or_documents=process_summary,
prompt_description=instructions,
examples=examples,
model_id="azure-openai",
batch_length=1,
max_workers=1,
max_char_buffer=100000,
prompt_validation_level=pv.PromptValidationLevel.OFF,
fence_output=False,
language_model_params={
"api_version": self.llm_config.AZURE_OPENAI_API_VERSION,
"temperature": self.llm_config.AZURE_OPENAI_TEMPERATURE,
"top_p": self.llm_config.AZURE_OPENAI_TOP_P,
"deployment_name": self.llm_config.AZURE_OPENAI_DEPLOYMENT_NAME,
"max_tokens": self.llm_config.AZURE_OPENAI_MAX_TOKENS,
"response_format": self.output_model,
},
)

My problem is that the resolver tries to align everything and not just the value of extraction_text. How can i specify what field the resolver should align or not align? Your help would be much appreciated. Thanks a lot in advance.

[h-vermeiren]

Hello,
Not sure if it answers your question, but I found that the json string that langextract expects in the infer step follows a different schema than the lx.data.Extraction() and that causes problems in the visualisation.
Basically, langextract expects the extraction_class value as key, and the extraction_text value as value, the attributes key needs to be prefixed with the extraction_class value.

{'Person': "John Doe", 'Person_attributes':{'first_name':"John"}}

This is how I fixed it:
a pydantic schema inspired by yours:

class Annotation(BaseModel):
    extraction_class: EntityType = Field(..., description="Entity type or class the entity belongs to")
    extraction_text: str = Field(..., description="Exact text from source that refers to the entity")
    attributes: Entity = Field(None, description="a valid object of one of the possible entity types")

class Extractions(BaseModel):
    extractions: List[Annotation]

    def to_lx_json(self) -> str:
        extractions=[]
        for ann in self.extractions:
            obj = ann.model_dump()
            newkey = obj.pop("extraction_class")
            newvalue = obj.pop("extraction_text")
            obj[newkey] = newvalue
            attributes = obj.pop("attributes")
            obj[f"{newkey}_attributes"] = attributes
            extractions.append(obj)
        return json.dumps({'extractions':extractions})

and then in infer:

# output_str = parsed.model_dump_json()    # old string for LangExtract
output_str = parsed.to_lx_json()      # new string for Langextract