60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
import argparse
|
|
from datetime import datetime
|
|
from typing import List, Optional
|
|
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_openai import ChatOpenAI
|
|
from pydantic import BaseModel, Field
|
|
|
|
# 1. Models
|
|
class PersonInfo(BaseModel):
|
|
name: str = Field(..., description="Имя человека")
|
|
age: Optional[int] = Field(None, description="Возраст (необязательно)")
|
|
profession: str = Field(..., description="Профессия")
|
|
skills: List[str] = Field(..., description="Список навыков")
|
|
|
|
class MeetingNotes(BaseModel):
|
|
date: datetime = Field(..., description="Дата встречи")
|
|
participants: List[str] = Field(..., description="Участники")
|
|
topics: List[str] = Field(..., description="Темы обсуждения")
|
|
decisions: List[str] = Field(..., description="Принятые решения")
|
|
next_steps: List[str] = Field(..., description="Следующие шаги")
|
|
|
|
# 2. Prompt and parser
|
|
prompt_template = """
|
|
Extract structured data from the following text.
|
|
Use the appropriate schema:
|
|
- PersonInfo for a person description
|
|
- MeetingNotes for meeting notes
|
|
Return only JSON matching the chosen schema.
|
|
Text: {text}
|
|
"""
|
|
prompt = PromptTemplate.from_template(prompt_template)
|
|
parser_person = PydanticOutputParser(pydantic_object=PersonInfo)
|
|
parser_meeting = PydanticOutputParser(pydantic_object=MeetingNotes)
|
|
|
|
# 3. Simple type guesser (very naive)
|
|
def choose_parser(text: str):
|
|
if "профессия" in text.lower() or "навыки" in text.lower():
|
|
return parser_person, PersonInfo
|
|
else:
|
|
return parser_meeting, MeetingNotes
|
|
|
|
# 4. Main chain
|
|
llm = ChatOpenAI(temperature=0)
|
|
|
|
def extract(text: str):
|
|
parser, model_cls = choose_parser(text)
|
|
chain = prompt | llm | parser
|
|
result = chain.invoke({"text": text})
|
|
return result, model_cls
|
|
|
|
# 5. CLI
|
|
if __name__ == "__main__":
|
|
ap = argparse.ArgumentParser(description="Extract structured data from raw text")
|
|
ap.add_argument("--text", required=True, help="Raw input text")
|
|
args = ap.parse_args()
|
|
out, model_cls = extract(args.text)
|
|
print(out.model_dump(indent=2))
|