101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
"""CLI tool for extracting structured data from text using LangChain and Pydantic.
|
|
|
|
The script supports two schemas:
|
|
- PersonInfo
|
|
- MeetingNotes
|
|
|
|
It automatically detects the type of input text and runs the appropriate chain.
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
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
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pydantic 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="Следующие шаги")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Prompt templates
|
|
# ---------------------------------------------------------------------------
|
|
PERSON_PROMPT = PromptTemplate(
|
|
input_variables=["text"],
|
|
template=(
|
|
"""
|
|
Извлечь из следующего текста информацию о человеке в формате JSON, соответствующем схеме PersonInfo.
|
|
Текст: {text}
|
|
JSON: """
|
|
),
|
|
)
|
|
|
|
MEETING_PROMPT = PromptTemplate(
|
|
input_variables=["text"],
|
|
template=(
|
|
"""
|
|
Извлечь из следующего текста информацию о встрече в формате JSON, соответствующем схеме MeetingNotes.
|
|
Текст: {text}
|
|
JSON: """
|
|
),
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper functions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def detect_schema(text: str) -> str:
|
|
"""Простейшая эвристика: если в тексте встречается слово "встреча" или "meeting" – считаем это встречей.
|
|
Иначе – человек.
|
|
"""
|
|
lowered = text.lower()
|
|
if "встреча" in lowered or "meeting" in lowered:
|
|
return "meeting"
|
|
return "person"
|
|
|
|
|
|
def build_chain(schema: str):
|
|
if schema == "person":
|
|
parser = PydanticOutputParser(pydantic_object=PersonInfo)
|
|
chain = PERSON_PROMPT | ChatOpenAI(temperature=0) | parser
|
|
else:
|
|
parser = PydanticOutputParser(pydantic_object=MeetingNotes)
|
|
chain = MEETING_PROMPT | ChatOpenAI(temperature=0) | parser
|
|
return chain
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Extract structured data from text.")
|
|
parser.add_argument("--text", required=True, help="Input text to parse")
|
|
args = parser.parse_args()
|
|
|
|
schema = detect_schema(args.text)
|
|
chain = build_chain(schema)
|
|
result = chain.invoke({"text": args.text})
|
|
|
|
# result is a Pydantic model instance
|
|
print(result.model_dump(indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|