77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
"""
|
||
Agent module that performs structured extraction using LangChain and Pydantic.
|
||
|
||
The module exposes a single public function `extract_structured_output` which:
|
||
1. Detects whether the input text describes a person or a meeting.
|
||
2. Builds an appropriate prompt with format instructions from the chosen parser.
|
||
3. Sends the request to the LLM via LangChain and returns the parsed Pydantic model instance.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from typing import Union, List
|
||
|
||
from langchain_core.output_parsers import PydanticOutputParser
|
||
from langchain_core.prompts import PromptTemplate
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
from models import PersonInfo, MeetingNotes
|
||
|
||
# LLM configuration – BroJS endpoint
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
||
api_key=None, # will be taken from env var JOURNAL_MCP_PAT
|
||
temperature=0.0,
|
||
)
|
||
|
||
# Prompt template – the same for both schemas; format instructions are injected.
|
||
PROMPT = PromptTemplate(
|
||
input_variables=["text", "format_instructions"],
|
||
template="""
|
||
You are a data extraction assistant.
|
||
Extract structured information from the following text and return it as JSON that matches the provided schema.
|
||
|
||
Text:
|
||
{text}
|
||
|
||
Schema format instructions:
|
||
{format_instructions}
|
||
|
||
Respond only with valid JSON. Do not add any extra text.
|
||
""",
|
||
)
|
||
|
||
def _detect_schema(text: str) -> type:
|
||
"""Heuristically determine whether the input describes a person or a meeting.
|
||
Returns the corresponding Pydantic model class.
|
||
"""
|
||
# Simple keyword checks – can be extended
|
||
if re.search(r"\bmeeting\b|\bdate\b|\bparticipants?\b", text, re.I):
|
||
return MeetingNotes
|
||
return PersonInfo
|
||
|
||
|
||
def extract_structured_output(text: str) -> Union[PersonInfo, MeetingNotes]:
|
||
"""Return a Pydantic model instance extracted from *text*.
|
||
|
||
Parameters
|
||
----------
|
||
text:
|
||
Free‑form input string.
|
||
|
||
Returns
|
||
-------
|
||
PersonInfo | MeetingNotes
|
||
Parsed data as a validated Pydantic object.
|
||
"""
|
||
schema = _detect_schema(text)
|
||
parser = PydanticOutputParser(pydantic_object=schema)
|
||
prompt = PROMPT.format(text=text, format_instructions=parser.get_format_instructions())
|
||
response = llm.invoke([prompt])
|
||
# The LLM returns a string; parse it
|
||
return parser.parse(response.content)
|
||
|
||
__all__ = ["extract_structured_output"]
|