Initial solution: add main.py
This commit is contained in:
@@ -0,0 +1,100 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user