Build a LangChain Customer Support Agent with Knowledge Base

import os
from typing import List
from langchain_community.document_loaders import (
    DirectoryLoader,
    UnstructuredMarkdownLoader,
    CSVLoader
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document

# Knowledge base structure
KB_DIRECTORY = "./knowledge_base"
VECTORSTORE_PATH = "./support_kb_vectorstore"

def load_knowledge_base(kb_dir: str) -> List[Document]:
    """Load all KB articles from a directory of markdown files."""
    loader = DirectoryLoader(
        kb_dir,
        glob="**/*.md",
        loader_cls=UnstructuredMarkdownLoader,
        show_progress=True
    )
    documents = loader.load()
    
    # Add FAQ CSV if it exists
    faq_path = os.path.join(kb_dir, "faqs.csv")
    if os.path.exists(faq_path):
        faq_loader = CSVLoader(
            file_path=faq_path,
            source_column="question",
            metadata_columns=["category", "priority"]
        )
        documents.extend(faq_loader.load())
    
    print(f"Loaded {len(documents)} documents from knowledge base")
    return documents

def ingest_knowledge_base(documents: List[Document]) -> Chroma:
    """Chunk, embed, and store KB documents in a vector store."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=600,
        chunk_overlap=80,
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    chunks = splitter.split_documents(documents)
    
    # Tag chunks with KB metadata
    for chunk in chunks:
        chunk.metadata.setdefault("source_type", "knowledge_base")
        chunk.metadata.setdefault("ingestion_date", "2026-05-31")
    
    print(f"Created {len(chunks)} chunks from {len(documents)} documents")
    
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=VECTORSTORE_PATH,
        collection_name="support_kb"
    )
    
    print(f"Stored {len(chunks)} chunks in vector store")
    return vectorstore

# Run ingestion
if os.path.exists(KB_DIRECTORY):
    docs = load_knowledge_base(KB_DIRECTORY)
    vectorstore = ingest_knowledge_base(docs)
else:
    # For demo: create in-memory store with sample data
    sample_docs = [
        Document(
            page_content="To reset your password: Click 'Forgot Password' on the login page, enter your email, check your inbox for the reset link, and follow the instructions. The link expires in 24 hours.",
            metadata={"source": "account_management.md", "category": "account", "priority": "high"}
        ),
        Document(
            page_content="Refund Policy: We offer full refunds within 30 days of purchase for unused subscriptions. To request a refund, contact billing@company.com with your order number. Processing takes 5-7 business days.",
            metadata={"source": "billing_policy.md", "category": "billing", "priority": "high"}
        ),
        Document(
            page_content="API rate limits: Free tier allows 100 requests/hour. Pro tier allows 10,000 requests/hour. Enterprise tier has custom limits. If you exceed your limit, you'll receive a 429 error.",
            metadata={"source": "api_documentation.md", "category": "technical", "priority": "medium"}
        ),
        Document(
            page_content="Data export: Go to Settings > Data Management > Export. Select date range and format (CSV, JSON, or XML). Large exports may take up to 30 minutes and will be emailed when ready.",
            metadata={"source": "data_management.md", "category": "features", "priority": "medium"}
        ),
    ]
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = Chroma.from_documents(
        documents=sample_docs,
        embedding=embeddings,
        collection_name="support_kb_demo"
    )
    print("Using demo knowledge base")

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnableLambda
from enum import Enum

class Intent(str, Enum):
    BILLING = "billing_question"
    TECHNICAL = "technical_issue"
    ACCOUNT = "account_management"
    FEATURE_REQUEST = "feature_request"
    COMPLAINT = "complaint"
    REFUND = "refund_request"
    GENERAL = "general_inquiry"
    ESCALATION = "escalation_request"

fast_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

CLASSIFICATION_PROMPT = ChatPromptTemplate.from_template("""
Classify this customer support message into exactly one category.

Categories:
- billing_question: questions about charges, invoices, payment methods
- technical_issue: bugs, errors, performance problems, API issues
- account_management: password reset, profile settings, access permissions
- feature_request: asking for new features or improvements
- complaint: expressing dissatisfaction or frustration
- refund_request: asking for money back
- general_inquiry: general questions not fitting other categories
- escalation_request: explicitly asking for a human agent or manager

Customer message: {message}

Respond with ONLY the category name, nothing else.
""")

classify_intent = CLASSIFICATION_PROMPT | fast_llm | StrOutputParser()

def get_intent(message: str) -> str:
    """Classify a user message and return the intent label."""
    result = classify_intent.invoke({"message": message})
    # Normalize output
    result = result.strip().lower()
    valid_intents = {i.value for i in Intent}
    return result if result in valid_intents else Intent.GENERAL.value

# Test classification
test_messages = [
    "I was charged twice for my subscription last month",
    "Getting a 429 error when calling the export endpoint",
    "I'd like to cancel and get a refund please",
    "Can I speak to a supervisor?",
    "How do I change my email address?",
]

for msg in test_messages:
    intent = get_intent(msg)
    print(f"Message: {msg[:50]}...")
    print(f"Intent: {intent}\n")

from langchain_core.retrievers import BaseRetriever
from langchain_core.documents import Document
from typing import List

INTENT_TO_CATEGORY = {
    "billing_question": "billing",
    "refund_request": "billing",
    "technical_issue": "technical",
    "account_management": "account",
    "feature_request": "features",
    "complaint": None,  # Search all categories
    "general_inquiry": None,
    "escalation_request": None,
}

def get_filtered_retriever(intent: str, k: int = 4):
    """Return a retriever filtered to the relevant KB category."""
    category = INTENT_TO_CATEGORY.get(intent)
    
    if category:
        search_kwargs = {
            "k": k,
            "filter": {"category": category}
        }
    else:
        search_kwargs = {"k": k}
    
    return vectorstore.as_retriever(search_kwargs=search_kwargs)

def retrieve_kb_articles(query: str, intent: str) -> List[Document]:
    """Retrieve relevant KB articles for a query given its intent."""
    retriever = get_filtered_retriever(intent)
    docs = retriever.invoke(query)
    
    if not docs:
        # Fall back to unfiltered search if filtered returns nothing
        fallback_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
        docs = fallback_retriever.invoke(query)
    
    return docs

def format_kb_context(docs: List[Document]) -> str:
    """Format retrieved documents into a context string."""
    if not docs:
        return "No relevant articles found in the knowledge base."
    
    formatted = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source", "knowledge_base")
        formatted.append(f"[Article {i} - {source}]\n{doc.page_content}")
    
    return "\n\n".join(formatted)

from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field
import json

class SupportResponse(BaseModel):
    answer: str = Field(description="The support response to give the customer")
    confidence: float = Field(description="Confidence score 0.0-1.0 that this answer fully resolves the issue")
    resolution_type: str = Field(description="'resolved', 'partial', or 'needs_escalation'")
    follow_up_question: str = Field(description="Optional clarifying question, empty string if not needed")

quality_llm = ChatOpenAI(model="gpt-4o", temperature=0.2)

RESPONSE_PROMPT = ChatPromptTemplate.from_template("""
You are a helpful, professional customer support agent.

Customer message: {message}
Customer intent: {intent}
Conversation history: {history}

Relevant knowledge base articles:
{context}

Instructions:
1. Answer the customer's question accurately using the KB articles provided
2. Be warm but concise — aim for 2-4 sentences
3. If the KB does not contain a clear answer, admit it and suggest next steps
4. Never make up policies, prices, or features not in the KB
5. If this requires account-specific information you don't have, say so

Respond in JSON format matching this schema:
{{
    "answer": "your response to the customer",
    "confidence": 0.0-1.0,
    "resolution_type": "resolved|partial|needs_escalation",
    "follow_up_question": "optional clarifying question or empty string"
}}
""")

def generate_response(
    message: str,
    intent: str,
    context: str,
    history: list
) -> SupportResponse:
    """Generate a support response with confidence scoring."""
    history_str = "\n".join([
        f"{msg['role'].upper()}: {msg['content']}"
        for msg in history[-6:]  # Last 3 turns
    ]) if history else "No previous messages"
    
    response_text = (RESPONSE_PROMPT | quality_llm | StrOutputParser()).invoke({
        "message": message,
        "intent": intent,
        "context": context,
        "history": history_str
    })
    
    # Parse JSON response
    try:
        data = json.loads(response_text)
        return SupportResponse(**data)
    except (json.JSONDecodeError, ValueError):
        # Fallback if JSON parsing fails
        return SupportResponse(
            answer=response_text,
            confidence=0.5,
            resolution_type="partial",
            follow_up_question=""
        )

from dataclasses import dataclass
from typing import Optional

ESCALATION_TRIGGERS = [
    "speak to a human",
    "talk to a person",
    "escalate",
    "supervisor",
    "manager",
    "this is unacceptable",
    "legal action",
    "lawyer",
    "complaint",
    "disgusting",
    "furious",
    "lawsuit"
]

@dataclass
class EscalationDecision:
    should_escalate: bool
    reason: str
    priority: str  # 'low', 'medium', 'high', 'urgent'
    suggested_team: str

def check_escalation(
    message: str,
    intent: str,
    response: SupportResponse,
    session_data: dict
) -> EscalationDecision:
    """Determine if this interaction should be escalated to a human agent."""
    
    message_lower = message.lower()
    
    # Explicit escalation request
    if intent == "escalation_request" or any(t in message_lower for t in ESCALATION_TRIGGERS):
        return EscalationDecision(
            should_escalate=True,
            reason="Customer explicitly requested human assistance",
            priority="high",
            suggested_team="general_support"
        )
    
    # Low confidence from LLM
    if response.confidence < 0.45:
        return EscalationDecision(
            should_escalate=True,
            reason=f"Low response confidence: {response.confidence:.2f}",
            priority="medium",
            suggested_team="technical_support" if intent == "technical_issue" else "general_support"
        )
    
    # LLM determined escalation needed
    if response.resolution_type == "needs_escalation":
        return EscalationDecision(
            should_escalate=True,
            reason="Agent determined issue requires human review",
            priority="medium",
            suggested_team="billing_team" if "billing" in intent else "general_support"
        )
    
    # Repeated contact about same issue
    if session_data.get("message_count", 0) > 4 and session_data.get("unresolved_turns", 0) > 2:
        return EscalationDecision(
            should_escalate=True,
            reason="Customer contacted support multiple times without resolution",
            priority="high",
            suggested_team="senior_support"
        )
    
    # Billing complaints above threshold
    if intent == "refund_request" and session_data.get("refund_amount", 0) > 500:
        return EscalationDecision(
            should_escalate=True,
            reason="High-value refund request requires approval",
            priority="high",
            suggested_team="billing_team"
        )
    
    return EscalationDecision(
        should_escalate=False,
        reason="Issue handled by automated agent",
        priority="low",
        suggested_team="none"
    )

from langchain_core.runnables import RunnableLambda
from typing import TypedDict, Optional
import datetime

class ConversationSession(TypedDict):
    session_id: str
    user_id: str
    history: list
    message_count: int
    unresolved_turns: int
    created_at: str

class AgentResponse(TypedDict):
    answer: str
    intent: str
    confidence: float
    escalated: bool
    escalation_reason: Optional[str]
    follow_up_question: str
    kb_sources: list

def run_support_agent(
    message: str,
    session: ConversationSession
) -> AgentResponse:
    """Run the full customer support agent pipeline."""
    
    # Step 1: Classify intent
    intent = get_intent(message)
    
    # Step 2: Retrieve KB articles
    kb_docs = retrieve_kb_articles(message, intent)
    context = format_kb_context(kb_docs)
    
    # Step 3: Generate response
    response = generate_response(
        message=message,
        intent=intent,
        context=context,
        history=session["history"]
    )
    
    # Step 4: Check escalation
    escalation = check_escalation(
        message=message,
        intent=intent,
        response=response,
        session_data={
            "message_count": session["message_count"],
            "unresolved_turns": session["unresolved_turns"]
        }
    )
    
    # Step 5: Update session history
    session["history"].append({"role": "user", "content": message})
    session["history"].append({"role": "assistant", "content": response.answer})
    session["message_count"] += 1
    
    if response.resolution_type != "resolved":
        session["unresolved_turns"] += 1
    else:
        session["unresolved_turns"] = 0
    
    # Step 6: Compose final response
    final_answer = response.answer
    
    if escalation.should_escalate:
        escalation_message = (
            f"\n\nI'm connecting you with a member of our {escalation.suggested_team.replace('_', ' ')} "
            f"team who can help you further. Expected wait time: 2-5 minutes."
        )
        final_answer += escalation_message
    
    if response.follow_up_question:
        final_answer += f"\n\n{response.follow_up_question}"
    
    return AgentResponse(
        answer=final_answer,
        intent=intent,
        confidence=response.confidence,
        escalated=escalation.should_escalate,
        escalation_reason=escalation.reason if escalation.should_escalate else None,
        follow_up_question=response.follow_up_question,
        kb_sources=[doc.metadata.get("source", "unknown") for doc in kb_docs]
    )

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uuid

app = FastAPI(title="Customer Support Agent API")

# In-memory session store (use Redis in production)
sessions = {}

class MessageRequest(BaseModel):
    message: str
    session_id: Optional[str] = None
    user_id: str = "anonymous"

@app.post("/support/chat")
async def chat(request: MessageRequest):
    # Get or create session
    session_id = request.session_id or str(uuid.uuid4())
    
    if session_id not in sessions:
        sessions[session_id] = ConversationSession(
            session_id=session_id,
            user_id=request.user_id,
            history=[],
            message_count=0,
            unresolved_turns=0,
            created_at=datetime.datetime.utcnow().isoformat()
        )
    
    session = sessions[session_id]
    
    try:
        response = run_support_agent(request.message, session)
        return {
            "session_id": session_id,
            "response": response["answer"],
            "intent": response["intent"],
            "confidence": response["confidence"],
            "escalated": response["escalated"],
            "escalation_reason": response["escalation_reason"],
            "kb_sources": response["kb_sources"]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Agent error: {str(e)}")

@app.get("/support/sessions/{session_id}")
async def get_session(session_id: str):
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    session = sessions[session_id]
    return {
        "session_id": session_id,
        "message_count": session["message_count"],
        "history_length": len(session["history"]),
        "unresolved_turns": session["unresolved_turns"]
    }

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

class ResolutionFeedback(BaseModel):
    session_id: str
    resolved: bool
    satisfaction_score: Optional[int] = None  # 1-5
    feedback_text: Optional[str] = None

feedback_store = []

@app.post("/support/feedback")
async def submit_feedback(feedback: ResolutionFeedback):
    feedback_store.append({
        "session_id": feedback.session_id,
        "resolved": feedback.resolved,
        "satisfaction_score": feedback.satisfaction_score,
        "feedback_text": feedback.feedback_text,
        "submitted_at": datetime.datetime.utcnow().isoformat()
    })
    return {"status": "feedback_recorded", "thank_you": True}

def compute_resolution_stats() -> dict:
    """Compute resolution statistics for the monitoring dashboard."""
    if not feedback_store:
        return {"total": 0, "resolution_rate": 0.0, "avg_satisfaction": 0.0}
    
    total = len(feedback_store)
    resolved = sum(1 for f in feedback_store if f["resolved"])
    scores = [f["satisfaction_score"] for f in feedback_store if f.get("satisfaction_score")]
    
    return {
        "total_sessions": total,
        "resolution_rate": resolved / total if total > 0 else 0.0,
        "avg_satisfaction": sum(scores) / len(scores) if scores else 0.0,
        "escalation_rate": sum(
            1 for sid in feedback_store
            if sessions.get(sid.get("session_id"), {}).get("escalated", False)
        ) / total if total > 0 else 0.0
    }

def run_test_conversation():
    """Simulate a test conversation with the support agent."""
    session = ConversationSession(
        session_id="test-001",
        user_id="test_user",
        history=[],
        message_count=0,
        unresolved_turns=0,
        created_at=datetime.datetime.utcnow().isoformat()
    )
    
    test_messages = [
        "Hi, I can't log into my account",
        "I tried resetting my password but the link doesn't work",
        "I need a refund, I've been charged twice",
        "Can I speak to a manager please"
    ]
    
    print("=" * 60)
    print("SUPPORT AGENT TEST CONVERSATION")
    print("=" * 60)
    
    for message in test_messages:
        print(f"\nUSER: {message}")
        response = run_support_agent(message, session)
        print(f"AGENT: {response['answer']}")
        print(f"Intent: {response['intent']} | Confidence: {response['confidence']:.2f}")
        if response['escalated']:
            print(f"ESCALATED: {response['escalation_reason']}")
        print("-" * 40)

run_test_conversation()