Освойте LangGraph: создание динамичного пайплайна для анализа текста с помощью ИИ

!pip install langgraph langchain langchain-openai python-dotenv

import os
from dotenv import load_dotenv
 
load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')

from langchain_openai import ChatOpenAI
 
llm = ChatOpenAI(model="gpt-4o-mini")
response = llm.invoke("Hello! Are you working?")
print(response.content)

from typing import TypedDict, List
 
class State(TypedDict):
    text: str
    classification: str
    entities: List[str]
    summary: str

def classification_node(state: State):
    prompt = PromptTemplate(
        input_variables=["text"],
        template="Классифицируйте следующий текст по категориям: Новости, Блог, Исследование или Другое.\n\nТекст:{text}\n\nКатегория:"
    )
    message = HumanMessage(content=prompt.format(text=state["text"]))
    classification = llm.invoke([message]).content.strip()
    return {"classification": classification}

def entity_extraction_node(state: State):
    prompt = PromptTemplate(
        input_variables=["text"],
        template="Извлеките все сущности (Персона, Организация, Локация) из текста. Результат – через запятую.\n\nТекст:{text}\n\nСущности:"
    )
    message = HumanMessage(content=prompt.format(text=state["text"]))
    entities = llm.invoke([message]).content.strip().split(", ")
    return {"entities": entities}

def summarization_node(state: State):
    prompt = PromptTemplate(
        input_variables=["text"],
        template="Сделайте краткое резюме текста в одном предложении.\n\nТекст:{text}\n\nРезюме:"
    )
    message = HumanMessage(content=prompt.format(text=state["text"]))
    summary = llm.invoke([message]).content.strip()
    return {"summary": summary}

workflow = StateGraph(State)
workflow.add_node("classification_node", classification_node)
workflow.add_node("entity_extraction", entity_extraction_node)
workflow.add_node("summarization", summarization_node)
workflow.set_entry_point("classification_node")
workflow.add_edge("classification_node", "entity_extraction")
workflow.add_edge("entity_extraction", "summarization")
workflow.add_edge("summarization", END)
app = workflow.compile()

sample_text = """ OpenAI объявила модель GPT-4, которая представляет собой крупную мультимодальную модель с человеческим уровнем производительности по разным профессиональным тестам. Она создана для улучшения согласованности и безопасности систем ИИ. Кроме того, модель более эффективна и масштабируема по сравнению с GPT-3. Ожидается, что GPT-4 будет доступна в ближайшие месяцы для исследований и разработок. """
state_input = {"text": sample_text}
result = app.invoke(state_input)
print("Классификация:", result["classification"])
print("\nСущности:", result["entities"])
print("\nРезюме:", result["summary"])

class EnhancedState(TypedDict):
    text: str
    classification: str
    entities: List[str]
    summary: str
    sentiment: str
 
def sentiment_node(state: EnhancedState):
    prompt = PromptTemplate(
        input_variables=["text"],
        template="Определите тональность текста: Позитивная, Негативная или Нейтральная?\n\nТекст:{text}\n\nТональность:"
    )
    message = HumanMessage(content=prompt.format(text=state["text"]))
    sentiment = llm.invoke([message]).content.strip()
    return {"sentiment": sentiment}
 
enhanced_workflow = StateGraph(EnhancedState)
enhanced_workflow.add_node("classification_node", classification_node)
enhanced_workflow.add_node("entity_extraction", entity_extraction_node)
enhanced_workflow.add_node("summarization", summarization_node)
enhanced_workflow.add_node("sentiment_analysis", sentiment_node)
enhanced_workflow.set_entry_point("classification_node")
enhanced_workflow.add_edge("classification_node", "entity_extraction")
enhanced_workflow.add_edge("entity_extraction", "summarization")
enhanced_workflow.add_edge("summarization", "sentiment_analysis")
enhanced_workflow.add_edge("sentiment_analysis", END)
enhanced_app = enhanced_workflow.compile()
 
enhanced_result = enhanced_app.invoke({"text": sample_text})
print("Классификация:", enhanced_result["classification"])
print("\nСущности:", enhanced_result["entities"])
print("\nРезюме:", enhanced_result["summary"])
print("\nТональность:", enhanced_result["sentiment"])

def route_after_classification(state: EnhancedState) -> str:
    category = state["classification"].lower()
    return category in ["news", "research"]
 
conditional_workflow = StateGraph(EnhancedState)
conditional_workflow.add_node("classification_node", classification_node)
conditional_workflow.add_node("entity_extraction", entity_extraction_node)
conditional_workflow.add_node("summarization", summarization_node)
conditional_workflow.add_node("sentiment_analysis", sentiment_node)
conditional_workflow.set_entry_point("classification_node")
conditional_workflow.add_conditional_edges("classification_node", route_after_classification, path_map={
    True: "entity_extraction",
    False: "summarization"
})
conditional_workflow.add_edge("entity_extraction", "summarization")
conditional_workflow.add_edge("summarization", "sentiment_analysis")
conditional_workflow.add_edge("sentiment_analysis", END)
conditional_app = conditional_workflow.compile()
 
# Тест с новостным текстом
news_text = """
OpenAI выпустила GPT-4 с улучшенной производительностью в академических и профессиональных задачах. Это значимый прорыв в согласованности и возможностях рассуждения.
"""
result = conditional_app.invoke({"text": news_text})
print("Классификация:", result["classification"])
print("Сущности:", result.get("entities", "Пропущено"))
print("Резюме:", result["summary"])
print("Тональность:", result["sentiment"])
 
# Тест с блогом
blog_text = """
Вот что я узнал после недели медитации в тишине. Без телефонов, без разговоров — только я, дыхание и глубокие осознания.
"""
result = conditional_app.invoke({"text": blog_text})
print("Классификация:", result["classification"])
print("Сущности:", result.get("entities", "Пропущено (не применимо)"))
print("Резюме:", result["summary"])
print("Тональность:", result["sentiment"])