Spaces:
Running
Running
import os | |
from dotenv import load_dotenv | |
from typing import TypedDict, List, Annotated | |
from langchain_groq import ChatGroq | |
from pydantic import BaseModel, Field | |
from langsmith import traceable | |
from langgraph.graph import StateGraph, START, END | |
from IPython.display import Image, display, Markdown | |
from langchain_core.messages import SystemMessage, HumanMessage | |
from langgraph.constants import Send | |
import operator | |
# Load environment variables | |
load_dotenv() | |
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY") | |
os.environ["LANGCHAIN_TRACING_V2"] = "true" | |
os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGCHAIN_API_KEY") | |
# Initialize LLM model | |
llm = ChatGroq(model="qwen-2.5-32b") | |
# Define Custom Data structures | |
class Topic(BaseModel): | |
"""Represents a learning topic with a name and description.""" | |
name: str = Field(description="Name of the learning topic.") | |
description: str = Field(description="Brief overview of the topic.") | |
class Topics(BaseModel): | |
"""Wrapper for a list of learning topics.""" | |
topics: List[Topic] = Field(description="List of topics to learn.") | |
# Augment the LLM with structured schema | |
planner = llm.with_structured_output(Topics) | |
# Define the state that carries data throughout the workflow | |
class State(TypedDict): | |
user_skills: str | |
user_goals: str | |
topics: List[Topic] | |
completed_topics: Annotated[List[str], operator.add] # Merging completed topics | |
learning_roadmap: str | |
# Worker state for topic processing | |
class WorkerState(TypedDict): | |
topic: Topic | |
completed_topics: List[str] | |
# Define Node Functions | |
def orchestrator(state: State): | |
"""Creates a study plan based on user skills and goals.""" | |
# LLM generates a structured study plan | |
study_plan = planner.invoke([ | |
SystemMessage( | |
content="Create a detailed study plan based on user skills and goals." | |
), | |
HumanMessage( | |
content=f"User skills: {state['user_skills']}\nUser goals: {state['user_goals']}" | |
), | |
]) | |
print("Study Plan:", study_plan) | |
return {"topics": study_plan.topics} # Returns generated topics | |
def llm_call(state: WorkerState): | |
"""Generates a content summary for a specific topic.""" | |
# LLM processes the topic and generates a summary | |
topic_summary = llm.invoke([ | |
SystemMessage( | |
content="Generate a content summary for the provided topic." | |
), | |
HumanMessage( | |
content=f"Topic: {state['topic'].name}\nDescription: {state['topic'].description}" | |
), | |
]) | |
return {"completed_topics": [topic_summary.content]} # Returns generated summary | |
def synthesizer(state: State): | |
"""Compiles topic summaries into a structured learning roadmap.""" | |
topic_summaries = state["completed_topics"] | |
learning_roadmap = "\n\n---\n\n".join(topic_summaries) # Formatting output | |
return {"learning_roadmap": learning_roadmap} # Returns final roadmap | |
# Define Conditional Edge Function | |
def assign_workers(state: State): | |
"""Assigns a worker (llm_call) to each topic in the plan.""" | |
return [Send("llm_call", {"topic": t}) for t in state["topics"]] # Creates worker tasks | |
# Build Workflow | |
learning_path_builder = StateGraph(State) | |
# Add nodes | |
learning_path_builder.add_node("orchestrator", orchestrator) | |
learning_path_builder.add_node("llm_call", llm_call) | |
learning_path_builder.add_node("synthesizer", synthesizer) | |
# Define execution order using edges | |
learning_path_builder.add_edge(START, "orchestrator") # Start with orchestrator | |
learning_path_builder.add_conditional_edges("orchestrator", assign_workers, ["llm_call"]) # Assign workers | |
learning_path_builder.add_edge("llm_call", "synthesizer") # Process topics | |
learning_path_builder.add_edge("synthesizer", END) # End workflow | |
# Compile workflow | |
learning_path_workflow = learning_path_builder.compile() | |
# ---------------------------- | |
# 5️⃣ Run the Workflow | |
# ---------------------------- | |
user_skills = "Python programming, basic machine learning concepts" | |
user_goals = "Learn advanced AI, master prompt engineering, and build AI applications" | |
state = learning_path_workflow.invoke( | |
{"user_skills": user_skills, "user_goals": user_goals} | |
) | |
# Display the final learning roadmap | |
Markdown(state["learning_roadmap"]) | |