LangGraph対DSPy：AIエージェントのための宣言型プロンプト最適化とグラフステートマシン

import dspy
from dspy.teleprompt import BootstrapFewShot
from typing import List, Dict, Any

# Configure DSPy with a local LLM (e.g., Ollama) or OpenAI
# For Ollama:
# llm = dspy.Ollama(model="llama3", max_tokens=2000)
# For OpenAI:
llm = dspy.OpenAI(model="gpt-4o-mini", max_tokens=2000)
dspy.settings.configure(lm=llm)

# Define a signature for generating search queries
class GenerateSearchQueries(dspy.Signature):
    """Generate a list of search queries for a given research question."""
    research_question: str = dspy.InputField(desc="The user's research question")
    search_queries: List[str] = dspy.OutputField(desc="A list of relevant search queries")

# Define a signature for summarizing search results
class SummarizeSearchResults(dspy.Signature):
    """Summarize a collection of search results into a concise overview."""
    search_results: List[str] = dspy.InputField(desc="A list of search result snippets")
    summary: str = dspy.OutputField(desc="A concise summary of the search results")

# Define a signature for identifying ambiguities or gaps
class IdentifyGaps(dspy.Signature):
    """Analyze a research summary and identify any ambiguities, missing information, or areas requiring clarification."""
    research_summary: str = dspy.InputField(desc="The current research summary")
    gaps_identified: str = dspy.OutputField(desc="A description of identified gaps or ambiguities, or 'None' if clear")
    requires_clarification: bool = dspy.OutputField(desc="True if user clarification is needed, False otherwise")

# Define DSPy Modules
class SearchQueryGenerator(dspy.Module):
    def __init__(self):
        super().__init__()
        self.generate_queries = dspy.Predict(GenerateSearchQueries)

    def forward(self, research_question: str) -> List[str]:
        prediction = self.generate_queries(research_question=research_question)
        return prediction.search_queries

class SearchResultSummarizer(dspy.Module):
    def __init__(self):
        super().__init__()
        self.summarize = dspy.Predict(SummarizeSearchResults)

    def forward(self, search_results: List[str]) -> str:
        prediction = self.summarize(search_results=search_results)
        return prediction.summary

class GapIdentifier(dspy.Module):
    def __init__(self):
        super().__init__()
        self.identify = dspy.Predict(IdentifyGaps)

    def forward(self, research_summary: str) -> Dict[str, Any]:
        prediction = self.identify(research_summary=research_summary)
        return {
            "gaps_identified": prediction.gaps_identified,
            "requires_clarification": prediction.requires_clarification
        }

# --- Optimization with BootstrapFewShot (example) ---
# In a real scenario, you'd have a dataset of (input, output) pairs
# For demonstration, we'll create a dummy dataset and compile.

# Dummy training data for GenerateSearchQueries
train_data_queries = [
    dspy.Example(research_question="Impact of AI on healthcare", search_queries=["AI in healthcare", "healthcare automation", "AI medical diagnostics"]),
    dspy.Example(research_question="Future of quantum computing", search_queries=["quantum computing trends", "quantum algorithms", "quantum hardware development"]),
]

# Dummy training data for SummarizeSearchResults
train_data_summaries = [
    dspy.Example(search_results=["Snippet 1 about AI", "Snippet 2 about healthcare"], summary="AI is transforming healthcare."),
    dspy.Example(search_results=["Snippet 1 about quantum", "Snippet 2 about future"], summary="Quantum computing holds future promise."),
]

# Dummy training data for IdentifyGaps
train_data_gaps = [
    dspy.Example(research_summary="AI is used in diagnostics.", gaps_identified="Does not specify types of AI or specific diagnostic applications.", requires_clarification=True),
    dspy.Example(research_summary="Quantum computing is a new field.", gaps_identified="None", requires_clarification=False),
]

# Define a simple metric for evaluation (e.g., checking if output is not empty)
def simple_metric(pred, gold, trace=None):
    return bool(pred.search_queries) if 'search_queries' in pred else bool(pred.summary) if 'summary' in pred else bool(pred.gaps_identified)

# Compile the modules
print("Compiling SearchQueryGenerator...")
teleprompter_queries = BootstrapFewShot(metric=simple_metric)
compiled_query_generator = teleprompter_queries.compile(SearchQueryGenerator(), trainset=train_data_queries)
print("SearchQueryGenerator compiled.")

print("Compiling SearchResultSummarizer...")
teleprompter_summaries = BootstrapFewShot(metric=simple_metric)
compiled_summarizer = teleprompter_summaries.compile(SearchResultSummarizer(), trainset=train_data_summaries)
print("SearchResultSummarizer compiled.")

print("Compiling GapIdentifier...")
teleprompter_gaps = BootstrapFewShot(metric=simple_metric)
compiled_gap_identifier = teleprompter_gaps.compile(GapIdentifier(), trainset=train_data_gaps)
print("GapIdentifier compiled.")

# Now, these compiled modules can be used within LangGraph
# For demonstration, let's test them:
# print("\nTesting compiled modules:")
# print(f"Queries: {compiled_query_generator.forward(research_question='Impact of 5G on IoT')}")
# print(f"Summary: {compiled_summarizer.forward(search_results=['5G enables faster IoT', 'IoT devices benefit from low latency'])}")
# print(f"Gaps: {compiled_gap_identifier.forward(research_summary='5G is fast.')}")

from typing import TypedDict, Annotated, List, Dict
import operator
from langgraph.graph import StateGraph, END
from langchain_community.tools import DuckDuckGoSearchRun # Example tool

# Define the state for our graph
class ResearchState(TypedDict):
    research_question: str
    search_queries: Annotated[List[str], operator.add]
    search_results: Annotated[List[str], operator.add]
    summary: str
    gaps_identified: str
    requires_clarification: bool
    user_clarification: str
    final_report: str
    iterations: int

# Initialize tools
search_tool = DuckDuckGoSearchRun()

# Define LangGraph nodes
def generate_queries_node(state: ResearchState) -> ResearchState:
    print("---GENERATING QUERIES---")
    question = state["research_question"]
    # Use the compiled DSPy module
    queries = compiled_query_generator.forward(research_question=question)
    return {"search_queries": queries, "iterations": state.get("iterations", 0) + 1}

def perform_search_node(state: ResearchState) -> ResearchState:
    print("---PERFORMING SEARCH---")
    queries = state["search_queries"]
    results = []
    for query in queries:
        print(f"Searching for: {query}")
        # In a real scenario, you'd handle rate limits, errors, etc.
        try:
            result = search_tool.run(query)
            results.append(result)
        except Exception as e:
            print(f"Search failed for '{query}': {e}")
    return {"search_results": results}

def summarize_results_node(state: ResearchState) -> ResearchState:
    print("---SUMMARIZING RESULTS---")
    results = state["search_results"]
    # Use the compiled DSPy module
    summary = compiled_summarizer.forward(search_results=results)
    return {"summary": summary}

def identify_gaps_node(state: ResearchState) -> ResearchState:
    print("---IDENTIFYING GAPS---")
    summary = state["summary"]
    # Use the compiled DSPy module
    gap_info = compiled_gap_identifier.forward(research_summary=summary)
    return {
        "gaps_identified": gap_info["gaps_identified"],
        "requires_clarification": gap_info["requires_clarification"]
    }

def human_in_the_loop_node(state: ResearchState) -> ResearchState:
    print("---HUMAN IN THE LOOP---")
    print(f"Current Summary: {state['summary']}")
    print(f"Gaps Identified: {state['gaps_identified']}")
    user_input = input("Clarification needed. Please provide additional context or guidance (type 'continue' to proceed without further input): ")
    return {"user_clarification": user_input}

def refine_report_node(state: ResearchState) -> ResearchState:
    print("---REFINING REPORT---")
    # This node would typically use another DSPy module for final report generation
    # For simplicity, we'll just combine existing info.
    final_report_content = (
        f"Research Question: {state['research_question']}\n\n"
        f"Summary of Findings:\n{state['summary']}\n\n"
    )
    if state['gaps_identified'] != 'None':
        final_report_content += f"Identified Gaps: {state['gaps_identified']}\n"
    if state['user_clarification']:
        final_report_content += f"User Clarification: {state['user_clarification']}\n"
    final_report_content += "\n--- END OF REPORT ---"
    return {"final_report": final_report_content}

# Define conditional edge logic
def decide_next_step(state: ResearchState) -> str:
    if state["requires_clarification"] and state.get("user_clarification") == None:
        print("---DECISION: CLARIFICATION NEEDED---")
        return "human_review"
    elif state["requires_clarification"] and state.get("user_clarification") != None and state["user_clarification"].lower() != 'continue':
        print("---DECISION: RE-EVALUATE AFTER CLARIFICATION---")
        # If user provided clarification, we might want to re-run search/summarize
        # For this example, we'll just proceed to refine, but in a real system,
        # you'd likely loop back to generate_queries or summarize_results.
        return "refine_report"
    else:
        print("---DECISION: PROCEED TO FINAL REPORT---")
        return "refine_report"

# Build the graph
workflow = StateGraph(ResearchState)

workflow.add_node("generate_queries", generate_queries_node)
workflow.add_node("perform_search", perform_search_node)
workflow.add_node("summarize_results", summarize_results_node)
workflow.add_node("identify_gaps", identify_gaps_node)
workflow.add_node("human_review", human_in_the_loop_node)
workflow.add_node("refine_report", refine_report_node)

workflow.set_entry_point("generate_queries")

workflow.add_edge("generate_queries", "perform_search")
workflow.add_edge("perform_search", "summarize_results")
workflow.add_edge("summarize_results", "identify_gaps")

# Conditional edge from identify_gaps
workflow.add_conditional_edges(
    "identify_gaps",
    decide_next_step,
    {
        "human_review": "human_review",
        "refine_report": "refine_report",
    },
)

# After human review, decide if we need to re-evaluate or finalize
workflow.add_conditional_edges(
    "human_review",
    decide_next_step, # Re-use the same decision logic
    {
        "human_review": "human_review", # Loop back if user didn't provide enough info (or if we want to ask again)
        "refine_report": "refine_report",
    },
)

workflow.add_edge("refine_report", END)

app = workflow.compile()

# Run the agent
initial_state = {"research_question": "What are the latest advancements in sustainable energy storage for grid applications?", "search_queries": [], "search_results": [], "summary": "", "gaps_identified": "", "requires_clarification": False, "user_clarification": None, "final_report": "", "iterations": 0}

print("\n--- STARTING RESEARCH AGENT ---")
for s in app.stream(initial_state):
    print(s)
    print("---")

print("\n--- FINAL REPORT ---")
final_state = app.invoke(initial_state)
print(final_state["final_report"])