""" 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()