52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
"""
|
|
Utilities for working with structured LLM outputs.
|
|
|
|
Provides helper functions to easily configure LLMs for structured output
|
|
across different providers (OpenAI, Anthropic, Google).
|
|
"""
|
|
|
|
from typing import Any, Dict, Type
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
def get_structured_llm(llm: Any, schema: Type[BaseModel]):
|
|
"""
|
|
Configure an LLM to return structured output based on a Pydantic schema.
|
|
|
|
Args:
|
|
llm: The LangChain LLM instance
|
|
schema: Pydantic BaseModel class defining the expected output structure
|
|
|
|
Returns:
|
|
Configured LLM that returns structured output
|
|
|
|
Example:
|
|
```python
|
|
from tradingagents.schemas import TradeDecision
|
|
from tradingagents.utils.structured_output import get_structured_llm
|
|
|
|
structured_llm = get_structured_llm(llm, TradeDecision)
|
|
response = structured_llm.invoke("Should I buy AAPL?")
|
|
# response is a dict matching TradeDecision schema
|
|
```
|
|
"""
|
|
return llm.with_structured_output(schema=schema.model_json_schema(), method="json_schema")
|
|
|
|
|
|
def extract_structured_response(response: Dict[str, Any], schema: Type[BaseModel]) -> BaseModel:
|
|
"""
|
|
Validate and parse a structured response into a Pydantic model.
|
|
|
|
Args:
|
|
response: Dictionary response from structured LLM
|
|
schema: Pydantic BaseModel class to validate against
|
|
|
|
Returns:
|
|
Validated Pydantic model instance
|
|
|
|
Raises:
|
|
ValidationError: If response doesn't match schema
|
|
"""
|
|
return schema(**response)
|