import uuid
from typing import TypedDict, NotRequired
from langchain.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware import ModelRequest, ModelResponse, AgentMiddleware
from langchain.messages import SystemMessage
from langgraph.checkpoint.memory import InMemorySaver
from typing import Callable
# Define skill structure
class Skill(TypedDict):
"""A skill that can be progressively disclosed to the agent."""
name: str
description: str
content: str
# Define skills with schemas and business logic
SKILLS: list[Skill] = [
{
"name": "sales_analytics",
"description": "Database schema and business logic for sales data analysis including customers, orders, and revenue.",
"content": """# Sales Analytics Schema
## Tables
### customers
- customer_id (PRIMARY KEY)
- name
- email
- signup_date
- status (active/inactive)
- customer_tier (bronze/silver/gold/platinum)
### orders
- order_id (PRIMARY KEY)
- customer_id (FOREIGN KEY -> customers)
- order_date
- status (pending/completed/cancelled/refunded)
- total_amount
- sales_region (north/south/east/west)
### order_items
- item_id (PRIMARY KEY)
- order_id (FOREIGN KEY -> orders)
- product_id
- quantity
- unit_price
- discount_percent
## Business Logic
**Active customers**: status = 'active' AND signup_date <= CURRENT_DATE - INTERVAL '90 days'
**Revenue calculation**: Only count orders with status = 'completed'. Use total_amount from orders table, which already accounts for discounts.
**Customer lifetime value (CLV)**: Sum of all completed order amounts for a customer.
**High-value orders**: Orders with total_amount > 1000
## Example Query
-- Get top 10 customers by revenue in the last quarter
SELECT
c.customer_id,
c.name,
c.customer_tier,
SUM(o.total_amount) as total_revenue
FROM customers c
JOIN orders o ON c.customer_id = o.customer_id
WHERE o.status = 'completed'
AND o.order_date >= CURRENT_DATE - INTERVAL '3 months'
GROUP BY c.customer_id, c.name, c.customer_tier
ORDER BY total_revenue DESC
LIMIT 10;
""",
},
{
"name": "inventory_management",
"description": "Database schema and business logic for inventory tracking including products, warehouses, and stock levels.",
"content": """# Inventory Management Schema
## Tables
### products
- product_id (PRIMARY KEY)
- product_name
- sku
- category
- unit_cost
- reorder_point (minimum stock level before reordering)
- discontinued (boolean)
### warehouses
- warehouse_id (PRIMARY KEY)
- warehouse_name
- location
- capacity
### inventory
- inventory_id (PRIMARY KEY)
- product_id (FOREIGN KEY -> products)
- warehouse_id (FOREIGN KEY -> warehouses)
- quantity_on_hand
- last_updated
### stock_movements
- movement_id (PRIMARY KEY)
- product_id (FOREIGN KEY -> products)
- warehouse_id (FOREIGN KEY -> warehouses)
- movement_type (inbound/outbound/transfer/adjustment)
- quantity (positive for inbound, negative for outbound)
- movement_date
- reference_number
## Business Logic
**Available stock**: quantity_on_hand from inventory table where quantity_on_hand > 0
**Products needing reorder**: Products where total quantity_on_hand across all warehouses is less than or equal to the product's reorder_point
**Active products only**: Exclude products where discontinued = true unless specifically analyzing discontinued items
**Stock valuation**: quantity_on_hand * unit_cost for each product
## Example Query
-- Find products below reorder point across all warehouses
SELECT
p.product_id,
p.product_name,
p.reorder_point,
SUM(i.quantity_on_hand) as total_stock,
p.unit_cost,
(p.reorder_point - SUM(i.quantity_on_hand)) as units_to_reorder
FROM products p
JOIN inventory i ON p.product_id = i.product_id
WHERE p.discontinued = false
GROUP BY p.product_id, p.product_name, p.reorder_point, p.unit_cost
HAVING SUM(i.quantity_on_hand) <= p.reorder_point
ORDER BY units_to_reorder DESC;
""",
},
]
# Create skill loading tool
@tool
def load_skill(skill_name: str) -> str:
"""Load the full content of a skill into the agent's context.
Use this when you need detailed information about how to handle a specific
type of request. This will provide you with comprehensive instructions,
policies, and guidelines for the skill area.
Args:
skill_name: The name of the skill to load (e.g., "sales_analytics", "inventory_management")
"""
# Find and return the requested skill
for skill in SKILLS:
if skill["name"] == skill_name:
return f"Loaded skill: {skill_name}\n\n{skill['content']}"
# Skill not found
available = ", ".join(s["name"] for s in SKILLS)
return f"Skill '{skill_name}' not found. Available skills: {available}"
# Create skill middleware
class SkillMiddleware(AgentMiddleware):
"""Middleware that injects skill descriptions into the system prompt."""
# Register the load_skill tool as a class variable
tools = [load_skill]
def __init__(self):
"""Initialize and generate the skills prompt from SKILLS."""
# Build skills prompt from the SKILLS list
skills_list = []
for skill in SKILLS:
skills_list.append(
f"- **{skill['name']}**: {skill['description']}"
)
self.skills_prompt = "\n".join(skills_list)
def wrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
"""Sync: Inject skill descriptions into system prompt."""
# Build the skills addendum
skills_addendum = (
f"\n\n## Available Skills\n\n{self.skills_prompt}\n\n"
"Use the load_skill tool when you need detailed information "
"about handling a specific type of request."
)
# Append to system message content blocks
new_content = list(request.system_message.content_blocks) + [
{"type": "text", "text": skills_addendum}
]
new_system_message = SystemMessage(content=new_content)
modified_request = request.override(system_message=new_system_message)
return handler(modified_request)
# Initialize your chat model (replace with your model)
# Example: from langchain_anthropic import ChatAnthropic
# model = ChatAnthropic(model="claude-3-5-sonnet-20241022")
from langchain_openai import ChatOpenAI
model = ChatOpenAI(model="gpt-4")
# Create the agent with skill support
agent = create_agent(
model,
system_prompt=(
"You are a SQL query assistant that helps users "
"write queries against business databases."
),
middleware=[SkillMiddleware()],
checkpointer=InMemorySaver(),
)
# Example usage
if __name__ == "__main__":
# Configuration for this conversation thread
thread_id = str(uuid.uuid4())
config = {"configurable": {"thread_id": thread_id}}
# Ask for a SQL query
result = agent.invoke(
{
"messages": [
{
"role": "user",
"content": (
"Write a SQL query to find all customers "
"who made orders over $1000 in the last month"
),
}
]
},
config
)
# Print the conversation
for message in result["messages"]:
if hasattr(message, 'pretty_print'):
message.pretty_print()
else:
print(f"{message.type}: {message.content}")
