haham7/nbc-chatbot
0
1"""2NBC Chat Agent - LangGraph-based conversation agent with guardrails.3 4Features:5- Multi-turn conversation history6- Off-topic guardrails7- Integration with NBCRetrievalPipeline (Qdrant + Neo4j)8- Reference document images in answers9"""10 11import os12import re13from dataclasses import dataclass, field14from enum import Enum15from typing import Optional, List, Dict, Any16 17from dotenv import load_dotenv18from pydantic import BaseModel19 20load_dotenv()21 22from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage23from langchain_core.outputs import ChatGeneration, ChatResult24from langgraph.graph import StateGraph, END25from langchain_openai import ChatOpenAI26 27# Import existing retrieval pipeline28from scripts.retrieval_pipeline import NBCRetrievalPipeline29 30 31@dataclass32class QueryCategory(Enum):33 """Query classification."""34 NBC_BUILDING = "building" # Structural, fire, plumbing, etc.35 NBC_ACCESSIBILITY = "accessibility"36 NBC_FIRE_SAFETY = "fire_safety"37 NBC_OCCUPANCY = "occupancy"38 NBC_FOUNDATION = "foundation"39 NBC_GENERAL = "general" # General NBC questions40 OFF_TOPIC = "off_topic"41 CLARIFICATION = "clarification"42 IMAGE_QUERY = "image_query"43 44 45@dataclass46class ConversationTurn:47 """Single conversation turn."""48 user_message: str49 assistant_message: str50 referenced_clauses: List[str]51 referenced_images: List[str]52 53 54@dataclass55class AgentState:56 """LangGraph state for conversation."""57 messages: List[BaseMessage] = field(default_factory=list)58 history: List[ConversationTurn] = field(default_factory=list)59 last_query: str = ""60 classification: QueryCategory = field(default_factory=lambda: QueryCategory.NBC_GENERAL)61 retrieved_context: str = ""62 referenced_clauses: List[str] = field(default_factory=list)63 referenced_images: List[str] = field(default_factory=list)64 is_off_topic: bool = False65 needs_clarification: bool = False66 67 68class NBCTopics:69 """Valid NBC 2016 topics for guardrails."""70 71 TOPICS = [72 "building",73 "structure",74 "fire",75 "safety",76 "plumbing",77 "foundation",78 "floor",79 "roof",80 "wall",81 "occupancy",82 "accessibility",83 "ramp",84 "stairs",85 "exit",86 "egress",87 "ventilation",88 "lighting",89 "natural light",90 "daylight",91 "toilet",92 "sanitary",93 "water",94 "drainage",95 "sewage",96 "lighting",97 "electrical",98 "air conditioning",99 "elevator",100 "lift",101 "parking",102 "setback",103 "ground coverage",104 "site coverage",105 "floor area",106 " FAR",107 "building height",108 "storey",109 "compartmentalization",110 "fire resistance",111 "wall",112 "column",113 "beam",114 "slab",115 " NBC ",116 "national building code",117 " part ",118 "section",119 "clause",120 "table",121 "figure",122 "annex",123 ]124 125 @classmethod126 def is_related(cls, query: str) -> bool:127 """Check if query is related to NBC."""128 query_lower = query.lower()129 return any(topic in query_lower for topic in cls.TOPICS)130 131 132class ConversationMemory:133 """Manages conversation history."""134 135 def __init__(self, max_turns: int = 10):136 self.max_turns = max_turns137 self.turns: List[ConversationTurn] = []138 139 def add_turn(self, turn: ConversationTurn) -> None:140 """Add a conversation turn."""141 self.turns.append(turn)142 if len(self.turns) > self.max_turns:143 self.turns = self.turns[-self.max_turns:]144 145 def get_context(self) -> str:146 """Get formatted context from history."""147 if not self.turns:148 return ""149 150 parts = ["=== Conversation History ==="]151 for i, turn in enumerate(self.turns[-5:], 1):152 parts.append(f"Q{i}: {turn.user_message}")153 parts.append(f"A{i}: {turn.assistant_message[:200]}...")154 if turn.referenced_clauses:155 parts.append(f"Clauses: {', '.join(turn.referenced_clauses)}")156 return "\n".join(parts)157 158 def get_referenced_clauses(self) -> List[str]:159 """Get all clauses referenced in conversation."""160 clauses = []161 for turn in self.turns:162 clauses.extend(turn.referenced_clauses)163 return list(set(clauses))164 165 166class QueryClassifier:167 """Classifies user queries."""168 169 KEYWORDS = {170 "accessibility": ["ramp", "lift", "elevator", "accessible", "disabled", "wheelchair"],171 "fire_safety": ["fire", "exit", "escape", "egress", "evacuation", "sprinkler", "alarm"],172 "occupancy": ["occupancy", "residential", "commercial", "industrial", "institutional"],173 "foundation": ["foundation", "soil", "bore", "pile", "raft"],174 }175 176 @classmethod177 def classify(cls, query: str) -> QueryCategory:178 """Classify query into NBC topic category."""179 query_lower = query.lower()180 181 for category, keywords in cls.KEYWORDS.items():182 if any(kw in query_lower for kw in keywords):183 return QueryCategory[category.upper()]184 185 if NBCTopics.is_related(query):186 return QueryCategory.NBC_GENERAL187 188 return QueryCategory.OFF_TOPIC189 190 191class NBCChatAgent:192 """193 LangGraph-based chat agent for NBC 2016.194 195 Flow:196 1. Classify query (on-topic / off-topic / clarification)197 2. Retrieve relevant context from Qdrant + Neo4j198 3. Generate response with clauses199 4. Store in memory200 201 Uses:202 - Qdrant: Hybrid search (dense + sparse vectors)203 - Neo4j: Cross-reference traversal204 - Reranking: BGE reranker205 """206 207 def __init__(208 self,209 retrieval_pipeline: NBCRetrievalPipeline,210 llm: Optional[ChatOpenAI] = None,211 max_history: int = 10,212 ):213 self.retrieval = retrieval_pipeline214 self.llm = llm or ChatOpenAI(215 model=os.getenv("CHAT_MODEL", "gpt-4o"),216 base_url=os.getenv("OPENAI_BASE_URL"),217 api_key=os.getenv("OPENAI_API_KEY"),218 temperature=0.3,219 )220 self.memory = ConversationMemory(max_turns=max_history)221 222 self._build_graph()223 224 def _build_graph(self) -> None:225 """Build LangGraph workflow."""226 workflow = StateGraph(AgentState)227 228 workflow.add_node("classify", self._classify_node)229 workflow.add_node("retrieve", self._retrieve_node)230 workflow.add_node("generate", self._generate_node)231 workflow.add_node("off_topic_response", self._off_topic_node)232 233 workflow.set_entry_point("classify")234 235 workflow.add_conditional_edges(236 "classify",237 self._should_retrieve,238 {239 "retrieve": "retrieve",240 "off_topic": "off_topic_response",241 "clarify": END,242 },243 )244 245 workflow.add_edge("retrieve", "generate")246 workflow.add_edge("generate", END)247 workflow.add_edge("off_topic_response", END)248 249 self.graph = workflow.compile()250 251 def _classify_node(self, state: AgentState) -> AgentState:252 """Classify the query."""253 last_msg = state.messages[-1].content if state.messages else ""254 state.last_query = last_msg255 256 classification = QueryClassifier.classify(last_msg)257 state.classification = classification258 state.is_off_topic = classification == QueryCategory.OFF_TOPIC259 260 if "?" in last_msg and len(last_msg.split()) < 5:261 state.needs_clarification = True262 263 return state264 265 def _should_retrieve(self, state: AgentState) -> str:266 """Determine next step based on classification."""267 if state.is_off_topic:268 return "off_topic"269 if state.needs_clarification:270 return "clarify"271 return "retrieve"272 273 def _retrieve_node(self, state: AgentState) -> AgentState:274 """Retrieve relevant context."""275 if not self.retrieval:276 return state277 278 try:279 context = self.retrieval.retrieve(state.last_query)280 281 context_parts = []282 clauses = []283 images = []284 285 for chunk in context.main_chunks[:3]:286 context_parts.append(287 f"[Clause {chunk.metadata.clause_number}] {chunk.text[:500]}"288 )289 clauses.append(chunk.metadata.clause_number)290 if chunk.metadata.image_path:291 images.append(chunk.metadata.image_path)292 293 for chunk in context.reference_chunks[:2]:294 context_parts.append(295 f"[Ref Clause {chunk.metadata.clause_number}] {chunk.text[:300]}"296 )297 clauses.append(chunk.metadata.clause_number)298 if chunk.metadata.image_path:299 images.append(chunk.metadata.image_path)300 301 state.retrieved_context = "\n\n".join(context_parts)302 state.referenced_clauses = list(set(clauses))303 state.referenced_images = images304 305 except Exception as e:306 state.retrieved_context = f"Error retrieving: {e}"307 308 return state309 310 def _generate_node(self, state: AgentState) -> AgentState:311 """Generate response using LLM."""312 history_context = self.memory.get_context()313 314 system_prompt = f"""You are an expert on the National Building Code of India 2016 (NBC 2016).315Answer questions based ONLY on the provided context. If information is not in the context, say so.316 317{history_context}318 319When answering:3201. Cite the clause number (e.g., "According to Clause 9.18.2...")3212. Reference figures/tables if relevant (e.g., "See Figure 9...")3223. Be specific about requirements, not vague323"""324 325 human_prompt = f"""Query: {state.last_query}326 327Context:328{state.retrieved_context}329 330Provide a specific answer citing the applicable NBC clauses."""331 332 try:333 response = self.llm.invoke(334 [SystemMessage(content=system_prompt), HumanMessage(content=human_prompt)]335 )336 assistant_message = response.content337 except Exception as e:338 assistant_message = f"I found relevant information but encountered an error generating the response: {e}"339 340 turn = ConversationTurn(341 user_message=state.last_query,342 assistant_message=assistant_message,343 referenced_clauses=state.referenced_clauses,344 referenced_images=state.referenced_images,345 )346 self.memory.add_turn(turn)347 348 state.messages.append(HumanMessage(content=state.last_query))349 state.messages.append(AIMessage(content=assistant_message))350 351 return state352 353 def _off_topic_node(self, state: AgentState) -> AgentState:354 """Generate off-topic response."""355 state.messages.append(356 AIMessage(357 content=(358 "I can only answer questions related to the "359 "National Building Code of India 2016 (NBC 2016). "360 "Please ask about building requirements, fire safety, "361 "accessibility, or other NBC-related topics."362 )363 )364 )365 return state366 367 def chat(self, user_message: str) -> str:368 """369 Process user message and return assistant response.370 371 Args:372 user_message: The user's query373 374 Returns:375 Assistant's response message376 """377 state = AgentState(378 messages=[HumanMessage(content=user_message)]379 )380 381 result = self.graph.invoke(state)382 383 response = result.messages[-1].content if result.messages else "I couldn't process that query."384 return response385 386 def get_referenced_content(self) -> Dict[str, Any]:387 """Get clauses and images referenced in conversation."""388 return {389 "clauses": self.memory.get_referenced_clauses(),390 "images": [391 img for turn in self.memory.turns[-3:]392 for img in turn.referenced_images393 ],394 }395 396 def clear_history(self) -> None:397 """Clear conversation history."""398 self.memory.turns = []399 400 401class Guardrails:402 """Guardrail functions for NBC chat."""403 404 @staticmethod405 def is_valid_query(query: str) -> bool:406 """Check if query is valid NBC-related question."""407 if len(query.strip()) < 5:408 return False409 if NBCTopics.is_related(query):410 return True411 return False412 413 @staticmethod414 def get_fallback_response() -> str:415 """Get fallback response for invalid queries."""416 return (417 "I specialize in answering questions about the "418 "National Building Code of India 2016. "419 "Please ask about building structural requirements, "420 "fire safety, accessibility, or other NBC-related topics."421 )422 423 424def create_agent(retrieval_pipeline, llm: Optional[ChatOpenAI] = None) -> NBCChatAgent:425 """Factory function to create chat agent."""426 return NBCChatAgent(427 retrieval_pipeline=retrieval_pipeline,428 llm=llm,429 max_history=10,430 )