add agent.py

agent.py

@@ -0,0 +1,76 @@
+"""
+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"]