232 lines
6.9 KiB
Python
232 lines
6.9 KiB
Python
"""
|
|
LangGraph with ChatOpenAI and Structured Output Example.
|
|
|
|
This demonstrates using LangGraph with ChatOpenAI's with_structured_output feature
|
|
for type-safe structured JSON responses.
|
|
"""
|
|
|
|
from typing import TypedDict, Literal, List
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langchain_openai import ChatOpenAI
|
|
from pydantic import BaseModel, Field
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
|
|
# ============ Helper to create model ============
|
|
|
|
def create_chat_model(model: str = "glm-4.5-air", temperature: float = 0.7):
|
|
"""Create a ChatOpenAI model instance."""
|
|
return ChatOpenAI(
|
|
model=model,
|
|
temperature=temperature,
|
|
max_retries=2,
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
base_url=os.getenv("OPENAI_BASE_URL", "https://open.bigmodel.cn/api/paas/v4"),
|
|
)
|
|
|
|
|
|
# ============ Structured Output Schemas ============
|
|
|
|
class SentimentAnalysis(BaseModel):
|
|
"""Schema for sentiment analysis output."""
|
|
sentiment: Literal["positive", "negative", "neutral"] = Field(
|
|
description="Overall sentiment of the text"
|
|
)
|
|
confidence: float = Field(
|
|
description="Confidence score between 0 and 1",
|
|
ge=0,
|
|
le=1
|
|
)
|
|
emotions: List[str] = Field(
|
|
description="List of detected emotions"
|
|
)
|
|
keywords: List[str] = Field(
|
|
description="Key words or phrases extracted"
|
|
)
|
|
analysis: str = Field(
|
|
description="Detailed analysis explanation"
|
|
)
|
|
|
|
|
|
class EntityExtraction(BaseModel):
|
|
"""Schema for entity extraction output."""
|
|
persons: List[str] = Field(default=[], description="Person names mentioned")
|
|
locations: List[str] = Field(default=[], description="Locations mentioned")
|
|
organizations: List[str] = Field(default=[], description="Organizations mentioned")
|
|
dates: List[str] = Field(default=[], description="Dates or time references")
|
|
|
|
|
|
class CombinedAnalysis(BaseModel):
|
|
"""Combined analysis result."""
|
|
sentiment: SentimentAnalysis
|
|
entities: EntityExtraction
|
|
summary: str = Field(description="Brief summary of the input text")
|
|
|
|
|
|
# ============ State Schema ============
|
|
|
|
class State(TypedDict):
|
|
"""State flows through the graph."""
|
|
input_text: str
|
|
sentiment_result: SentimentAnalysis | None
|
|
entity_result: EntityExtraction | None
|
|
final_summary: str
|
|
|
|
|
|
# ============ Node Functions ============
|
|
|
|
def analyze_sentiment(state: State) -> dict:
|
|
"""Analyze sentiment using with_structured_output."""
|
|
print(f"Analyzing sentiment for: '{state['input_text']}'...")
|
|
|
|
# Create model with structured output
|
|
model = create_chat_model().with_structured_output(SentimentAnalysis)
|
|
|
|
prompt = f"""
|
|
你是一个情感分析专家。请分析用户输入的情感。
|
|
|
|
输入文本：{state['input_text']}
|
|
|
|
请按照指定的 JSON Schema 返回分析结果，包括 sentiment, confidence, emotions, keywords 和 analysis。
|
|
"""
|
|
|
|
result: SentimentAnalysis = model.invoke(prompt)
|
|
|
|
print(f" Sentiment: {result.sentiment}")
|
|
print(f" Confidence: {result.confidence:.0%}")
|
|
print(f" Emotions: {', '.join(result.emotions)}")
|
|
|
|
return {"sentiment_result": result}
|
|
|
|
|
|
def extract_entities(state: State) -> dict:
|
|
"""Extract entities using with_structured_output."""
|
|
print("Extracting entities...")
|
|
|
|
# Create model with structured output
|
|
model = create_chat_model().with_structured_output(EntityExtraction)
|
|
|
|
prompt = f"""
|
|
你是一个命名实体识别专家。请从文本中提取实体。
|
|
|
|
输入文本：{state['input_text']}
|
|
|
|
请提取 persons, locations, organizations, dates 等实体信息。
|
|
"""
|
|
|
|
result: EntityExtraction = model.invoke(prompt)
|
|
|
|
print(f" Persons: {result.persons}")
|
|
print(f" Locations: {result.locations}")
|
|
print(f" Organizations: {result.organizations}")
|
|
|
|
return {"entity_result": result}
|
|
|
|
|
|
def generate_summary(state: State) -> dict:
|
|
"""Generate a combined summary using with_structured_output."""
|
|
print("Generating final summary...")
|
|
|
|
sentiment = state.get("sentiment_result")
|
|
entities = state.get("entity_result")
|
|
|
|
# Create model with structured output
|
|
model = create_chat_model().with_structured_output(CombinedAnalysis)
|
|
|
|
prompt = f"""
|
|
请对以下分析结果生成综合摘要：
|
|
|
|
输入文本：{state['input_text']}
|
|
|
|
情感分析结果：
|
|
- 情感：{sentiment.sentiment if sentiment else 'N/A'}
|
|
- 置信度：{sentiment.confidence if sentiment else 'N/A'}
|
|
- 情绪：{sentiment.emotions if sentiment else 'N/A'}
|
|
- 关键词：{sentiment.keywords if sentiment else 'N/A'}
|
|
|
|
实体提取结果：
|
|
- 人名：{entities.persons if entities else 'N/A'}
|
|
- 地点：{entities.locations if entities else 'N/A'}
|
|
- 组织：{entities.organizations if entities else 'N/A'}
|
|
- 日期：{entities.dates if entities else 'N/A'}
|
|
|
|
请生成一个简短的摘要。
|
|
"""
|
|
|
|
result: CombinedAnalysis = model.invoke(prompt)
|
|
|
|
return {"final_summary": result.summary}
|
|
|
|
|
|
# ============ Build Graph ============
|
|
|
|
def create_analysis_graph():
|
|
"""Create and compile the analysis workflow graph."""
|
|
|
|
builder = StateGraph(State)
|
|
|
|
# Add nodes
|
|
builder.add_node("sentiment", analyze_sentiment)
|
|
builder.add_node("entities", extract_entities)
|
|
builder.add_node("summarize", generate_summary)
|
|
|
|
# Add edges
|
|
builder.add_edge(START, "sentiment")
|
|
builder.add_edge("sentiment", "entities")
|
|
builder.add_edge("entities", "summarize")
|
|
builder.add_edge("summarize", END)
|
|
|
|
return builder.compile()
|
|
|
|
|
|
# ============ Main ============
|
|
|
|
def main():
|
|
"""Run the structured analysis agent."""
|
|
graph = create_analysis_graph()
|
|
|
|
print("=== LangGraph with ChatOpenAI Structured Output ===\n")
|
|
|
|
# Test with a sample text
|
|
test_text = "今天天气真好，我和朋友们一起去公园野餐，大家都很开心！"
|
|
|
|
print(f"Input text: {test_text}\n")
|
|
print("-" * 50)
|
|
|
|
# Run the graph
|
|
initial_state = {
|
|
"input_text": test_text,
|
|
"sentiment_result": None,
|
|
"entity_result": None,
|
|
"final_summary": ""
|
|
}
|
|
|
|
result = graph.invoke(initial_state)
|
|
|
|
print("-" * 50)
|
|
print("\n=== Final Summary ===")
|
|
print(result["final_summary"])
|
|
|
|
# Access structured pydantic models
|
|
print("\n=== Structured Sentiment Data ===")
|
|
sentiment: SentimentAnalysis = result["sentiment_result"]
|
|
print(f" sentiment: {sentiment.sentiment}")
|
|
print(f" confidence: {sentiment.confidence}")
|
|
print(f" emotions: {sentiment.emotions}")
|
|
print(f" keywords: {sentiment.keywords}")
|
|
print(f" analysis: {sentiment.analysis}")
|
|
|
|
print("\n=== Structured Entity Data ===")
|
|
entities: EntityExtraction = result["entity_result"]
|
|
print(f" persons: {entities.persons}")
|
|
print(f" locations: {entities.locations}")
|
|
print(f" organizations: {entities.organizations}")
|
|
print(f" dates: {entities.dates}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|