import argparse import json import os from langchain_core.prompts import PromptTemplate from agent import build_agent, create_demo_database from rag_agent import build_rag_agent from compressor import get_llm def format_rag_citations(rag_res): """Extracts metadata from the RAG context and appends it as formatted citations.""" ans = rag_res.get("answer", "") context = rag_res.get("context", []) if not context: return ans sources = set() for doc in context: source_val = doc.metadata.get("source", "Unknown Document") basename = os.path.basename(source_val) page = doc.metadata.get("page") if page is not None: # PDF pages are 0-indexed in Langchain sources.add(f"{basename} (Page {page + 1})") else: sources.add(basename) if sources: ans += "\n\n**Sources:**\n" + "\n".join([f"- `{s}`" for s in sources]) return ans def create_router(): """Builds a simple LLM router to classify user intent as structured (SQL), unstructured (RAG), both, or general.""" llm = get_llm() prompt = PromptTemplate.from_template( "You are an intelligent router for a company assistant. " "Your job is to classify the user's question into one of four categories:\n" "1. 'sql' - The question requires querying a relational database (e.g. sales, revenue, users, customers, products, orders).\n" "2. 'rag' - The question requires looking up unstructured documents (e.g. policies, office locations, rules, guidelines, text).\n" "3. 'both' - The question bridges both domains (e.g. 'Which country generated the most revenue and what is our refund policy there?').\n" "4. 'general' - A general conversational question or greeting (e.g. 'Hi', 'Who are you?').\n" "Output ONLY a valid JSON object with a single key 'route' and value 'sql', 'rag', 'both', or 'general'. Do not output any other text or markdown.\n\n" "Question: {question}" ) # This chain will output a JSON string return prompt | llm def init_system(db_path="data/demo.sqlite"): if db_path == "data/demo.sqlite" and not os.path.exists("data/demo.sqlite"): os.makedirs("data", exist_ok=True) create_demo_database("data/demo.sqlite") sql_agent, _ = build_agent(db_path) rag_agent = build_rag_agent() router = create_router() return sql_agent, rag_agent, router def ask_multi_agent(q, sql_agent, rag_agent, router, chat_history_str=""): # Inject chat history into the Router's classification prompt so it understands follow-ups router_q = q if chat_history_str: router_q = f"Previous Chat History:\n{chat_history_str}\n\nLatest Question: {q}" try: route_res = router.invoke({"question": router_q}).content route_res = route_res.replace('```json', '').replace('```', '').strip() route = json.loads(route_res).get("route", "rag") except Exception: route = "rag" route = str(route).lower().strip() # Inject chat history into the Agents' prompts so they can answer follow-ups correctly agent_q = q if chat_history_str: agent_q = f"Previous Chat History:\n{chat_history_str}\n\nLatest Question: {q}" if route == "general": try: llm = get_llm() res = llm.invoke(agent_q) return route, res.content except Exception as e: return "error", f"General LLM failed: {e}" elif route == "both": sql_ans = "Failed to fetch SQL data." try: sql_res = sql_agent.invoke({"input": agent_q}) sql_ans = sql_res["output"] except Exception as e: sql_ans = f"SQL Agent error: {e}" # Create a targeted query for the RAG agent using the SQL context! llm = get_llm() rag_query_prompt = PromptTemplate.from_template( "User's multi-part question: {question}\n" "We already queried the SQL database and found this: {sql_answer}\n\n" "Your task: Formulate a targeted search query to look up the remaining unanswered parts of the user's question in the company documents. " "If the remaining part depends on the SQL answer (e.g. a specific product), include that context. " "Output ONLY the search query, without any quotes or preamble." ) rag_q = (rag_query_prompt | llm).invoke({"question": agent_q, "sql_answer": sql_ans}).content.strip() rag_ans = "Failed to fetch RAG data." rag_citations = "" try: rag_res = rag_agent.invoke({"input": rag_q}) rag_ans = rag_res["answer"] # Extract citations for the Combiner to append at the end context = rag_res.get("context", []) sources = set() for doc in context: source_val = doc.metadata.get("source", "Unknown Document") basename = os.path.basename(source_val) page = doc.metadata.get("page") if page is not None: sources.add(f"{basename} (Page {page + 1})") else: sources.add(basename) if sources: rag_citations = "\n\n**Sources:**\n" + "\n".join([f"- `{s}`" for s in sources]) except Exception as e: rag_ans = f"RAG Agent error: {e}" combiner_prompt = PromptTemplate.from_template( "You are an intelligent synthesis assistant. You have received answers from two expert systems regarding the user's question.\n\n" "User's Question: {question}\n\n" "SQL Agent (Database) Answer:\n{sql_answer}\n\n" "RAG Agent (Documents) Answer:\n{rag_answer}\n\n" "Please synthesize these into a single, cohesive, and helpful response." ) chain = combiner_prompt | llm try: final_ans = chain.invoke({"question": agent_q, "sql_answer": sql_ans, "rag_answer": rag_ans}).content # Append citations to the final synthesized answer return route, final_ans + rag_citations except Exception as e: return "error", f"Combiner LLM failed: {e}" elif route == "sql": try: res = sql_agent.invoke({"input": agent_q}) return route, res["output"] except Exception as e: try: res = rag_agent.invoke({"input": agent_q}) return "rag (fallback from sql)", format_rag_citations(res) except Exception as e2: return "error", f"Both agents failed. SQL Error: {e}, RAG Error: {e2}" else: try: res = rag_agent.invoke({"input": agent_q}) return route, format_rag_citations(res) except Exception as e: try: res = sql_agent.invoke({"input": agent_q}) return "sql (fallback from rag)", res["output"] except Exception as e2: return "error", f"Both agents failed. RAG Error: {e}, SQL Error: {e2}" def main(): parser = argparse.ArgumentParser(description="Multi-Agent Master Router") parser.add_argument("--db", default="data/demo.sqlite", help="SQLite database path") parser.add_argument("--question", help="Question to ask (omit for interactive mode)") args = parser.parse_args() print("Loading Multi-Agent System...") try: sql_agent, rag_agent, router = init_system(args.db) except Exception as e: print(f"Error initializing system: {e}") return print("\n[OK] Multi-Agent System Ready!") def process_query(q): route, ans = ask_multi_agent(q, sql_agent, rag_agent, router) print(f" [Router Decision: Sending query to the {route.upper()} Agent]") return ans if args.question: print(f"\nUser: {args.question}") answer = process_query(args.question) print(f"\nAgent: {answer}") else: print("\nMulti-Agent Chat ready. Ask about sales, users, or company policies! Type 'quit' to exit.") while True: try: q = input("\nYou: ").strip() except (KeyboardInterrupt, EOFError): break if q.lower() in ["quit", "exit", "q"]: break if not q: continue answer = process_query(q) print(f"\nAgent: {answer}") if __name__ == "__main__": main()