51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
import os
|
|
import json
|
|
from typing import List
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.output_parsers import JsonOutputParser
|
|
from langchain_core.prompts import PromptTemplate
|
|
from pydantic import BaseModel
|
|
|
|
from models import TaskCard
|
|
|
|
class TaskParser:
|
|
"""Parser that converts raw task descriptions into TaskCard objects using a LLM."""
|
|
|
|
def __init__(self) -> None:
|
|
self.llm = ChatOpenAI(
|
|
model="openai/gpt-oss-20b:free",
|
|
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
|
api_key=os.getenv("JOURNAL_MCP_PAT"),
|
|
temperature=0.1,
|
|
)
|
|
self.parser = JsonOutputParser(pydantic_object=TaskCard)
|
|
self.prompt = PromptTemplate(
|
|
input_variables=["raw_text"],
|
|
template="""
|
|
Parse the following raw task description into a JSON object matching the TaskCard schema.
|
|
The output must be valid JSON and contain all required fields.
|
|
If a field cannot be determined, use null for optional fields or an empty list for lists.
|
|
|
|
Raw text:
|
|
{raw_text}
|
|
|
|
JSON output:
|
|
""",
|
|
)
|
|
|
|
def _run_llm(self, raw_text: str) -> TaskCard:
|
|
chain = self.prompt | self.llm | self.parser
|
|
return chain.invoke({"raw_text": raw_text})
|
|
|
|
def parse(self, raw_text: str) -> TaskCard:
|
|
return self._run_llm(raw_text)
|
|
|
|
def batch_parse(self, texts: List[str]) -> List[TaskCard]:
|
|
return [self.parse(t) for t in texts]
|
|
|
|
def save_to_file(self, card: TaskCard, filename: str) -> None:
|
|
with open(filename, "w", encoding="utf-8") as f:
|
|
json.dump(card.dict(), f, ensure_ascii=False, indent=2)
|
|
|
|
# End of parser.py
|