import asyncio import time # from functools import lru_cache from typing import Dict, Any, Literal, Optional, List from fastapi import HTTPException from langchain_core.messages import HumanMessage, AIMessage, ToolMessage, SystemMessage from langchain_core.tools import StructuredTool from langgraph.graph import StateGraph, END, START, MessagesState from langgraph.types import Command from langgraph.prebuilt import ToolNode from langgraph.prebuilt import create_react_agent # Internal imports from ..llm_utils import LLM_UTILS from ..logger import log from .chatschemas import ( RetrievalArgs, AgentProductResponse, AgentRequest, QueryArgs, RAGArgs, ) from .product_tools import ( get_variants, get_orders_by_email, get_order_details_by_id, get_active_coupons, ) from .chatutilities import ( _retrieval, _query_products, get_chat_history, _fetch_products_by_ids, _get_product_samples, _get_products, get_customer_id_by_email, knowledge_retriever, ) from .prompts import ( get_router_system_prompt, get_tool_model_system_prompt, get_supervisor_prompt, _get_create_order_agent_prompt, get_rag_agent_system_prompt, ) from ..AnalyzerAgent.user_facts_analyzer import get_saved_facts # ------------------- # Tools # ------------------- retrieval_tool = StructuredTool.from_function( func=_retrieval, args_schema=RetrievalArgs, name="retrieve_products", description="Retrieve products from shop catalog via natural language query.", ) product_query_tool = StructuredTool.from_function( func=_query_products, args_schema=QueryArgs, name="find_products", description="Find products matching via a mongo query filter (e.g. category, price, brand) with optional sort and limit", ) # rag_tool = StructuredTool.from_function( # func=metadata_retriever, # args_schema=RAGArgs, # name="retrieve_knowledge_entries", # description="Retrieve knowledge base entries relevant to a query using Qdrant embeddings, Returns the most relevant entries with their similarity scores.", # ) # ------------------- # MultiTenantSalesAgent # ------------------- class MultiTenantSalesAgent: def __init__(self, chat_id: str, store_id: Optional[str]): self.chat_id = chat_id self.store_id = store_id self.user_creds = None # Bind tools (works regardless of store_id) self.model_with_tools = LLM_UTILS.llm.bind_tools( [retrieval_tool, product_query_tool] ) # Structured output LLM self.sales_llm_router = LLM_UTILS.llm.with_structured_output( AgentProductResponse ) # Safely fetch product samples and titles only if store_id is valid if self.store_id: self.product_samples = _get_product_samples(self.store_id) self.product_titles = _get_products(self.store_id) else: self.product_samples = [] self.product_titles = [] def _create_workflow(self): graph = StateGraph(MessagesState) graph.add_node("agent", self._call_model) graph.add_node("tools", ToolNode([retrieval_tool, product_query_tool])) graph.add_edge(START, "agent") graph.add_conditional_edges("agent", self._should_continue, ["tools", END]) graph.add_edge("tools", "agent") return graph.compile(name="recommendation_agent") def _should_continue(self, state: MessagesState): last = state["messages"][-1] return "tools" if getattr(last, "tool_calls", None) else END def _call_model(self, state: MessagesState): start = time.time() last_message = state["messages"][-1] if isinstance(last_message, ToolMessage): prompt_to_use = get_router_system_prompt(self.product_titles) response = self.sales_llm_router.invoke([prompt_to_use] + state["messages"]) ai_message = AIMessage( content=[ {"answer": response.answer}, {"products": response.product_ids}, ] ) log.info(f"_call_model(general) finished in {time.time() - start:.2f}s") else: prompt_to_use = get_tool_model_system_prompt( self.product_samples, self.product_titles ) response = self.model_with_tools.invoke([prompt_to_use] + state["messages"]) ai_message = response log.info( f"_call_model(tool_calling) finished in {time.time() - start:.2f}s" ) return {"messages": [ai_message]} def _get_supervisor_router( self, state: MessagesState ) -> Command[ Literal[ "orderagent", "recommendation_agent", "humanagent", "rag_agent", "__end__" ] ]: """ Routes queries to orderagent, RecommendationAgent. """ router_sytstem_prompt = get_supervisor_prompt() # Use LLM router to classify the query result = LLM_UTILS.shopify_supervisor_router.invoke( [router_sytstem_prompt] + state["messages"] ) classification = getattr(result, "classification", "").lower() general = getattr(result, "general", "") update = None log.info(f"triage outcome is classified: {classification}, general: {general}") # Map classification to LangGraph Command if classification.lower() == "orderagent": log.info(f"Routing to orderagent") goto = "orderagent" elif classification.lower() == "humanagent": log.info(f"Routing to humanagent") goto = "humanagent" update = {"messages": [AIMessage(content=general)]} elif classification.lower() == "recommendation_agent": log.info(f"Routing to recommendation_agent") goto = "recommendation_agent" elif classification.lower() == "rag_agent": log.info(f"Routing to rag_agent") goto = "rag_agent" else: log.info(f"Routing to end with general response") goto = END update = {"messages": [AIMessage(content=general)]} return Command(goto=goto, update=update) def _create_order_agent(self): # user_facts = _customer_context_cache.get(self.chat_id, {}).copy() user_facts = get_saved_facts(self.chat_id) user_email = self.user_creds.get("user_email") if user_email: self.user_creds["user_customer_id"] = ( get_customer_id_by_email(user_email, self.store_id) or None ) else: self.user_creds["user_customer_id"] = None prompt = _get_create_order_agent_prompt(self.user_creds, user_facts) orderagent = create_react_agent( model=LLM_UTILS.llm, # substitute your language model (e.g., ChatOpenAI or Anthropic) tools=[ get_orders_by_email, get_order_details_by_id, get_variants, get_active_coupons, ], prompt=prompt, name="orderagent", ) return orderagent def orderagent(self, state: MessagesState) -> Command[Literal["__end__"]]: log.info(f"orderagent invoked") orderagent = self._create_order_agent() log.info(f"orderagent retrieved") result = orderagent.invoke(state) return Command( update={"messages": [AIMessage(content=result["messages"][-1].content)]}, # We want our workers to ALWAYS "report back" to the supervisor when done goto=END, ) def humanagent(self, state: MessagesState) -> Command[Literal["__end__"]]: last_message = state["messages"][-1] # if isinstance(last_message, ToolMessage): # answer = "Transferring to human agent. Please wait..." # else: # answer = last_message.content log.info(f"Human Agent invoked {last_message}") return Command( update={ "messages": [ AIMessage( content=[ { "True": "Please wait while I connect you to a human agent." }, {"False": "Do not call human agent over this."}, ] ) ] }, # We want our workers to ALWAYS "report back" to the supervisor when done goto=END, ) def _create_supervisor_agent(self): workflow = StateGraph(MessagesState) workflow.add_node(self._get_supervisor_router) workflow.add_node("recommendation_agent", self.recommendation_agent) workflow.add_node("rag_agent", self.rag_agent) workflow.add_node("orderagent", self.orderagent) workflow.add_node("humanagent", self.humanagent) workflow.add_edge(START, "_get_supervisor_router") return workflow.compile() def _get_agent_executor(self): return self._create_workflow() def recommendation_agent(self, state: MessagesState) -> Command[Literal["__end__"]]: recommendation_agent = self._get_agent_executor() result = recommendation_agent.invoke(state) log.info(f"Recommendation Agent invoked") return Command( update={"messages": [AIMessage(content=result["messages"][-1].content)]}, goto=END, ) def _create_rag_agent(self): rag_system_prompt = get_rag_agent_system_prompt() rag_agent = create_react_agent( model=LLM_UTILS.llm, tools=[knowledge_retriever], prompt=rag_system_prompt ) return rag_agent def rag_agent(self, state: MessagesState) -> Command[Literal["__end__"]]: log.info("Invoking RAG agent...") rag_agent = self._create_rag_agent() result = rag_agent.invoke(state) try: # If result is a BaseMessage or a dict-like object if hasattr(result, "content"): content = result.content elif isinstance(result, dict) and "content" in result: content = result["content"] elif isinstance(result, dict) and "messages" in result: content = result["messages"][-1].content else: # Fallback: try string conversion content = str(result) except Exception as e: log.error(f"Error extracting content from RAG agent result: {e}") content = "Error: could not extract RAG response content." log.info("RAG Agent invoked successfully.") return Command( update={"messages": [AIMessage(content=content)]}, goto=END, ) async def chat_with_ai(self, data: AgentRequest, user_id:str) -> Dict[str, Any]: try: if not all([data.text, data.chat_id]): raise HTTPException(status_code=400, detail="Missing required fields") self.user_creds = { "username": data.user_name, "user_email": data.user_email, } history = get_chat_history(data.chat_id) messages = [ ( HumanMessage(content=m["content"]) if m["role"] == "human" else AIMessage(content=m["content"]) ) for m in history ] + [HumanMessage(content=data.text)] executor = self._get_agent_executor() loop = asyncio.get_event_loop() result = await loop.run_in_executor( None, lambda: executor.invoke( {"messages": messages}, config={ "configurable": { "store_id": self.store_id, "user_id": user_id, "chat_id": data.chat_id, } }, ), ) last = result["messages"][-1] if isinstance(last.content, list): answer, products = last.content[0].get("answer"), last.content[1].get( "products", [] ) else: answer, products = last.content, [] return { "data": { "chat_id": data.chat_id, "ai_response": { "shopifyRecommendation": { "products": _fetch_products_by_ids(products) }, "woocommerceRecommendations": None, "chatresponseanswer": answer, }, } } except Exception as e: log.exception("chat_with_ai failed") raise HTTPException(status_code=500, detail="Internal server error") def build_chat_messages(self, chat_id: str, user_text: str) -> List: """ Fetch chat history and return a list of messages with the current user input appended. Args: chat_id (str): The chat session ID. user_text (str): The current user message. Returns: List of HumanMessage and AIMessage objects. """ history = get_chat_history(chat_id) messages = [ ( HumanMessage(content=m["content"]) if m["role"] == "user" else AIMessage(content=m["content"]) ) for m in history ] # Append current user message messages.append(HumanMessage(content=user_text)) return messages async def chat_with_ai_supervisor(self, data: AgentRequest, user_id:str): try: # Validate required fields if not all([data.text, data.chat_id]): raise HTTPException(status_code=400, detail="Missing required fields") # Prepare user credentials and chat history self.user_creds = { "username": data.user_name, "user_email": data.user_email, } # history = get_chat_history(data.chat_id) messages = self.build_chat_messages(data.chat_id, data.text) # Create supervisor agent and configuration supervisor_app = self._create_supervisor_agent() config = { "configurable": {"store_id": self.store_id, "user_id": user_id}, **(data.shopify_creds or {}), **(self.user_creds or {}), } # Run supervisor safely in background thread result = await asyncio.to_thread( supervisor_app.invoke, {"messages": messages}, config=config ) # log.info(f"Whole Agent respone {result}") last_msg = ( result.get("messages", [])[-1] if result.get("messages") else None ) if not last_msg: raise HTTPException( status_code=500, detail="No response from supervisor agent" ) content = last_msg.content answer, products = None, [] # Extract structured data if content is list if isinstance(content, list): answer = next( (item.get("answer") for item in content if "answer" in item), None ) products = next( (item.get("products") for item in content if "products" in item), [] ) # ✅ Final rule: # If content is NOT a list AND no "answer" key was found → return full content if not isinstance(content, list) and answer is None: chatresponseanswer = content else: chatresponseanswer = answer # Build product responses products_response = [] if products: fetched = _fetch_products_by_ids(products) or [] products_response = [p for p in fetched if p] # Extract human escalation messages humananswer = ( content if isinstance(content, list) and any("True" in i or "False" in i for i in content) else [] ) # Return structured response return { "data": { "chat_id": data.chat_id, "ai_response": { "shopifyRecommendation": {"products": products_response}, "woocommerceRecommendations": None, "chatresponseanswer": chatresponseanswer, "humanAgent": humananswer, }, } } except HTTPException: raise # keep explicit HTTP errors except Exception as e: log.exception("chat_with_ai_supervisor failed: %s", e) raise HTTPException(status_code=500, detail="Internal server error")