add parser.py
This commit is contained in:
@@ -0,0 +1,50 @@
|
||||
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
|
||||
Reference in New Issue
Block a user