Structured Output

class JSONSchema:
    schema: Mapping[str, Any]
    name: str
    description: Optional[str] = None
    strict: Optional[bool] = False
    
    def __init__(
        self,
        schema: Mapping[str, Any],
        name: str,
        description: Optional[str] = None,
        strict: Optional[bool] = False
    ):
        """
        JSON schema that structured output must adhere to.
        
        Parameters:
        - schema: JSON schema that structured output must adhere to
        - name: Name of the schema
        - description: Description of the schema
        - strict: Whether the schema should be strictly enforced
        """

    def to_json(self) -> Mapping[str, Any]:
        """Convert schema to JSON format for API requests."""

    @classmethod
    def from_pydantic(cls, model_class: Type[BaseModel]) -> JSONSchema:
        """
        Create a JSONSchema from a Pydantic model class.
        
        Parameters:
        - model_class: A Pydantic BaseModel class
        
        Returns:
        JSONSchema instance with schema generated from the Pydantic model
        
        Raises:
        ValueError: If the provided class is not a Pydantic BaseModel
        """

ResponseFormat = Union[JSONSchema, Type[BaseModel]]

from aleph_alpha_client import JSONSchema, ChatRequest, Client
from pydantic import BaseModel
from typing import List

client = Client(token="your-api-token")

# Manual JSON schema
schema = JSONSchema(
    schema={
        'type': 'object',
        'title': 'PersonInfo',
        'properties': {
            'name': {
                'type': 'string',
                'title': 'Name'
            },
            'age': {
                'type': 'integer',
                'title': 'Age'
            },
            'skills': {
                'type': 'array',
                'items': {'type': 'string'},
                'title': 'Skills'
            }
        },
        'required': ['name', 'age']
    },
    name="person_info",
    description="Extract person information",
    strict=True
)

# Use in chat request
request = ChatRequest(
    model="pharia-1-llm-7b-control",
    messages=[{
        "role": "user",
        "content": "Extract the person's information: John is 30 years old and knows Python and JavaScript"
    }],
    response_format=schema
)

response = client.chat(request)
print(response.message.content)  # Structured JSON output

# Using Pydantic models (automatic schema generation)
class PersonInfo(BaseModel):
    """Person information extraction model."""
    name: str
    age: int
    skills: List[str] = []

# Create schema from Pydantic model
pydantic_schema = JSONSchema.from_pydantic(PersonInfo)

request = ChatRequest(
    model="pharia-1-llm-7b-control",
    messages=[{
        "role": "user", 
        "content": "Extract info: Sarah is 25 and specializes in data science and machine learning"
    }],
    response_format=pydantic_schema
)

response = client.chat(request)

# Parse structured response
person_data = PersonInfo.model_validate_json(response.message.content)
print(f"Name: {person_data.name}, Age: {person_data.age}")
print(f"Skills: {', '.join(person_data.skills)}")

# Direct Pydantic model usage
request = ChatRequest(
    model="pharia-1-llm-7b-control",
    messages=[{
        "role": "user",
        "content": "Analyze this product review and extract key information"
    }],
    response_format=PersonInfo  # Direct Pydantic model
)