Initial commit: add LangGraph examples
This commit is contained in:
Eric Wang
231
chatopenai_structured.py
Normal file
231
chatopenai_structured.py
Normal file
@@ -0,0 +1,231 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user