74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
"""Base output parser class."""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional
|
|
|
|
from llama_index.bridge.pydantic import Field
|
|
from llama_index.core.query_pipeline.query_component import (
|
|
ChainableMixin,
|
|
InputKeys,
|
|
OutputKeys,
|
|
QueryComponent,
|
|
validate_and_convert_stringable,
|
|
)
|
|
from llama_index.types import BaseOutputParser
|
|
|
|
|
|
@dataclass
|
|
class StructuredOutput:
|
|
"""Structured output class."""
|
|
|
|
raw_output: str
|
|
parsed_output: Optional[Any] = None
|
|
|
|
|
|
class OutputParserException(Exception):
|
|
pass
|
|
|
|
|
|
class ChainableOutputParser(BaseOutputParser, ChainableMixin):
|
|
"""Chainable output parser."""
|
|
|
|
# TODO: consolidate with base at some point if possible.
|
|
|
|
def _as_query_component(self, **kwargs: Any) -> QueryComponent:
|
|
"""Get query component."""
|
|
return OutputParserComponent(output_parser=self)
|
|
|
|
|
|
class OutputParserComponent(QueryComponent):
|
|
"""Output parser component."""
|
|
|
|
output_parser: BaseOutputParser = Field(..., description="Output parser.")
|
|
|
|
class Config:
|
|
arbitrary_types_allowed = True
|
|
|
|
def _run_component(self, **kwargs: Any) -> Dict[str, Any]:
|
|
"""Run component."""
|
|
output = self.output_parser.parse(kwargs["input"])
|
|
return {"output": output}
|
|
|
|
async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
|
|
"""Run component."""
|
|
# NOTE: no native async for output parser
|
|
return self._run_component(**kwargs)
|
|
|
|
def _validate_component_inputs(self, input: Any) -> Any:
|
|
"""Validate component inputs during run_component."""
|
|
input["input"] = validate_and_convert_stringable(input["input"])
|
|
return input
|
|
|
|
def set_callback_manager(self, callback_manager: Any) -> None:
|
|
"""Set callback manager."""
|
|
|
|
@property
|
|
def input_keys(self) -> Any:
|
|
"""Input keys."""
|
|
return InputKeys.from_keys({"input"})
|
|
|
|
@property
|
|
def output_keys(self) -> Any:
|
|
"""Output keys."""
|
|
return OutputKeys.from_keys({"output"})
|