import logging import re from typing import List, Dict, Any, Optional from langchain_core.prompts import ChatPromptTemplate from pydantic import BaseModel, Field from .llm_client import get_llm from .neo4j_client import Neo4jClient from .vector_retriever import VectorRetriever logger = logging.getLogger(__name__) # Patterns that indicate multi-hop queries MULTIHOP_PATTERNS = [ r"who manages.*who", r"who leads.*that", r"what depends on.*that", r"what.*connected to.*through", r"who reports to.*who", r"which.*works at.*that", r"what.*built by.*that", r"find.*path between", r"how.*related to", r"who.*manages the team", r"what.*the manager of", r"list all.*connected", ] class EntityPair(BaseModel): """Two entities to find a path between.""" entity_a: str = Field(description="The first entity name") entity_b: str = Field(description="The second entity name") class QueryVariations(BaseModel): """Exactly 2 distinct search query variations of the input question.""" variations: List[str] = Field(description="List containing exactly 2 distinct search query variations of the input question") class MultiHopReasoner: def __init__(self): logger.info("Initializing MultiHopReasoner service.") self.neo4j_client = Neo4jClient() self.vector_retriever = VectorRetriever() self.llm = get_llm(temperature=0.0) # Entity extraction prompt for multi-hop queries self.entity_extract_prompt = ChatPromptTemplate.from_messages([ ("system", ( "Extract the two main entities from this multi-hop question. " "Return them as entity_a and entity_b." )), ("human", "{question}") ]) self.entity_extract_chain = self.entity_extract_prompt | self.llm.with_structured_output(EntityPair) # Query variations prompt and chain for query expansion self.variations_prompt = ChatPromptTemplate.from_messages([ ("system", ( "You are an expert search query variation assistant. " "Generate exactly 2 distinct search query variations of the input question to find connected entities in a graph." )), ("human", "{question}") ]) self.variations_chain = self.variations_prompt | self.llm.with_structured_output(QueryVariations) # Prompt instruction to summarize path connections self.prompt = ChatPromptTemplate.from_messages([ ("system", ( "You are an AI analyst specialized in explaining graph connections.\n" "You will be given a path of nodes and relationships from a knowledge graph showing how two entities are connected.\n" "Your job is to summarize this connection path in a clear, natural paragraph.\n\n" "Format rules:\n" "- Clearly mention each step of the connection.\n" "- Keep the explanation factual based *only* on the provided path info." )), ("human", ( "Explain the connection between '{entity_a}' and '{entity_b}' based on this graph path:\n\n" "{path_details}" )) ]) self.chain = self.prompt | self.llm @staticmethod def is_multihop_query(question: str) -> bool: """Detect if a question requires multi-hop reasoning based on pattern matching.""" q = question.lower() return any(re.search(pat, q) for pat in MULTIHOP_PATTERNS) def extract_entities_from_query(self, question: str) -> Optional[Dict[str, str]]: """Use LLM to extract entity pair from a multi-hop question.""" try: result: EntityPair = self.entity_extract_chain.invoke({"question": question}) return {"entity_a": result.entity_a.strip(), "entity_b": result.entity_b.strip()} except Exception as e: logger.error("Failed to extract entities from multi-hop query: %s", str(e)) return None def _get_nodes_and_rels(self, path_obj): if isinstance(path_obj, list): nodes = [path_obj[i] for i in range(0, len(path_obj), 2)] relationships = [path_obj[i] for i in range(1, len(path_obj), 2)] return nodes, relationships else: return list(path_obj.nodes), list(path_obj.relationships) def _get_rel_type(self, rel): if isinstance(rel, str): return rel if isinstance(rel, dict): return rel.get("type", "RELATED_TO") if hasattr(rel, "type"): return rel.type return "RELATED_TO" def _resolve_hybrid_entity_name(self, name: str, user_id: str) -> str: """Resolves an entity name to the closest existing entity name in the graph.""" if not name: return name name_clean = name.strip() # 1. First, check if exact match exists exact_cypher = ( "MATCH (e:Entity {name: $name, user_id: $user_id}) " "RETURN e.name AS name LIMIT 1" ) records = self.neo4j_client.execute_query(exact_cypher, {"name": name_clean, "user_id": str(user_id)}) if records: return records[0]["name"] # 2. Case-insensitive exact match case_cypher = ( "MATCH (e:Entity {user_id: $user_id}) " "WHERE toLower(e.name) = toLower($name) " "RETURN e.name AS name LIMIT 1" ) records = self.neo4j_client.execute_query(case_cypher, {"name": name_clean, "user_id": str(user_id)}) if records: return records[0]["name"] # 3. CONTAINS / partial match in Cypher contains_cypher = ( "MATCH (e:Entity {user_id: $user_id}) " "WHERE toLower(e.name) CONTAINS toLower($name) OR toLower($name) CONTAINS toLower(e.name) " "RETURN e.name AS name LIMIT 1" ) records = self.neo4j_client.execute_query(contains_cypher, {"name": name_clean, "user_id": str(user_id)}) if records: return records[0]["name"] # 4. Fetch all entity names and use difflib to find the closest match all_cypher = ( "MATCH (e:Entity {user_id: $user_id}) " "RETURN DISTINCT e.name AS name" ) try: all_records = self.neo4j_client.execute_query(all_cypher, {"user_id": str(user_id)}) all_names = [r["name"] for r in all_records if r.get("name")] import difflib matches = difflib.get_close_matches(name_clean, all_names, n=1, cutoff=0.3) if matches: logger.info("Fuzzy resolved entity '%s' -> '%s'", name_clean, matches[0]) return matches[0] except Exception as e: logger.error("Failed fuzzy matching name resolution: %s", str(e)) return name_clean def find_alternative_paths(self, entity_a: str, entity_b: str, user_id: str, max_paths: int = 5) -> List[List[Dict]]: """Find multiple alternative paths of different lengths between two entities.""" entity_a = self._resolve_hybrid_entity_name(entity_a, user_id) entity_b = self._resolve_hybrid_entity_name(entity_b, user_id) cypher = ( "MATCH p = (a:Entity {name: $entity_a, user_id: $user_id})-[*1..4]-(b:Entity {name: $entity_b, user_id: $user_id}) " "RETURN p LIMIT 30" ) params = { "entity_a": entity_a, "entity_b": entity_b, "user_id": str(user_id) } try: records = self.neo4j_client.execute_query(cypher, params) all_paths = [] for record in records: path_obj = record.get("p") if not path_obj: continue nodes, relationships = self._get_nodes_and_rels(path_obj) # Filter out paths with cycles (node repetition) node_names = [dict(n).get("name", "Unknown") for n in nodes] if len(node_names) != len(set(node_names)): continue path_steps = [] for i in range(len(relationships)): now_name = dict(nodes[i]).get("name", "Unknown") next_name = dict(nodes[i + 1]).get("name", "Unknown") rel = relationships[i] rel_type = self._get_rel_type(rel) source_doc = "" page = 1 if isinstance(rel, dict): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) elif hasattr(rel, "get"): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) else: try: source_doc = dict(rel).get("source_doc", "") page = dict(rel).get("page", 1) except Exception: pass chunk_text = "" if source_doc and page: chunk_text = self.vector_retriever.get_chunk_by_page(source_doc, page, user_id) or "" path_steps.append({ "source": now_name, "target": next_name, "type": rel_type, "source_doc": source_doc, "page": page, "chunk_text": chunk_text }) if path_steps: all_paths.append(path_steps) return all_paths except Exception as e: logger.error("Failed to find alternative paths: %s", str(e)) return [] def explain_connection(self, entity_a: str, entity_b: str, user_id: str) -> Dict[str, Any]: """ Finds the shortest path between two entities in Neo4j and uses the LLM to explain the connection. """ # Resolve names using hybrid fuzzy matching entity_a = self._resolve_hybrid_entity_name(entity_a, user_id) entity_b = self._resolve_hybrid_entity_name(entity_b, user_id) logger.info("Finding connection between '%s' and '%s' (User: %s)", entity_a, entity_b, user_id) # 1. Fetch shortest path from Neo4j cypher = ( "MATCH p = shortestPath(" " (a:Entity {name: $entity_a, user_id: $user_id})-[*..5]-(b:Entity {name: $entity_b, user_id: $user_id})" ") " "RETURN p" ) params = { "entity_a": entity_a, "entity_b": entity_b, "user_id": str(user_id) } try: records = self.neo4j_client.execute_query(cypher, params) if not records or not records[0].get("p"): logger.info("No connection path found between '%s' and '%s'.", entity_a, entity_b) return { "found": False, "explanation": f"No indirect connection (up to 4 hops) was found between '{entity_a}' and '{entity_b}' in the knowledge graph.", "path": [] } path_obj = records[0]["p"] nodes, relationships = self._get_nodes_and_rels(path_obj) # 2. Extract path details for prompt serialization path_steps = [] serialized_path = [] for i in range(len(relationships)): node_now = nodes[i] node_next = nodes[i + 1] rel = relationships[i] now_name = dict(node_now).get("name", "Unknown") next_name = dict(node_next).get("name", "Unknown") now_type = dict(node_now).get("type", "Entity") next_type = dict(node_next).get("type", "Entity") rel_type = self._get_rel_type(rel) # Extract source_doc and page source_doc = "" page = 1 if isinstance(rel, dict): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) elif hasattr(rel, "get"): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) else: try: source_doc = dict(rel).get("source_doc", "") page = dict(rel).get("page", 1) except Exception: pass # Get chunk text from vector retriever chunk_text = "" if source_doc and page: chunk_text = self.vector_retriever.get_chunk_by_page(source_doc, page, user_id) or "" step_str = f"({now_name} [{now_type}]) --[{rel_type}]--> ({next_name} [{next_type}])" path_steps.append(step_str) # Keep tracking representation for the response payload serialized_path.append({ "source": now_name, "source_type": now_type, "target": next_name, "target_type": next_type, "type": rel_type, "source_doc": source_doc, "page": page, "chunk_text": chunk_text }) path_details = "\n".join(path_steps) logger.info("Found path with %d hops: %s", len(relationships), path_details) # 3. Ask LLM to summarize/explain this path response = self.chain.invoke({ "entity_a": entity_a, "entity_b": entity_b, "path_details": path_details }) explanation = response.content.strip() # Find alternative paths alt_paths = self.find_alternative_paths(entity_a, entity_b, user_id) # Filter out the main path main_path_key = tuple((s["source"], s["target"], s["type"]) for s in serialized_path) alternative_paths = [ p for p in alt_paths if tuple((s["source"], s["target"], s["type"]) for s in p) != main_path_key ] # Generate explanation for alternative paths alternative_paths_with_explanations = [] for alt in alternative_paths[:2]: alt_steps = [] for step in alt: alt_steps.append(f"({step['source']}) --[{step['type']}]--> ({step['target']})") alt_details = "\n".join(alt_steps) try: alt_response = self.chain.invoke({ "entity_a": entity_a, "entity_b": entity_b, "path_details": alt_details }) alt_explanation = alt_response.content.strip() except Exception as alt_err: logger.error("Failed to generate explanation for alt path: %s", str(alt_err)) alt_explanation = "Could not generate reasoning explanation for this alternative path." alternative_paths_with_explanations.append({ "hops": alt, "explanation": alt_explanation }) return { "found": True, "explanation": explanation, "path": serialized_path, "alternative_paths": alternative_paths_with_explanations, "hop_count": len(relationships) } except Exception as e: logger.error("Failed to execute path reasoning query: %s", str(e), exc_info=True) return { "found": False, "explanation": "An internal error occurred while analyzing the connection.", "path": [] } def generate_query_variations(self, question: str) -> List[str]: """Generate 2 alternative variations of the query using the LLM with structured schema enforcement.""" try: result: QueryVariations = self.variations_chain.invoke({"question": question}) return [v.strip() for v in result.variations if v.strip()] except Exception as e: logger.error("Failed to generate query variations: %s", str(e)) return [] def find_all_reasoning_paths(self, query: str, primary_a: str, primary_b: str, user_id: str) -> Dict[str, Any]: """ Runs query expansion to find multiple alternative connecting paths between entities related to the query. Returns the primary path with generated explanation, and other unique paths with empty explanations (for on-demand generation). """ pairs_to_search = [] if primary_a and primary_b: resolved_a = self._resolve_hybrid_entity_name(primary_a, user_id) resolved_b = self._resolve_hybrid_entity_name(primary_b, user_id) pairs_to_search.append((resolved_a, resolved_b)) # Generate query variations using LLM if a natural language query is provided if query: variations = self.generate_query_variations(query) logger.info("Generated query variations: %s", variations) for var in variations: pair = self.extract_entities_from_query(var) if pair: resolved_var_a = self._resolve_hybrid_entity_name(pair["entity_a"], user_id) resolved_var_b = self._resolve_hybrid_entity_name(pair["entity_b"], user_id) pair_tuple = (resolved_var_a, resolved_var_b) if pair_tuple not in pairs_to_search: pairs_to_search.append(pair_tuple) all_unique_paths = [] seen_signatures = set() # Helper to generate path signature def get_signature(hops): return "-".join(f"{h.get('source') or h.get('from')}->{h.get('type') or h.get('rel')}->{h.get('target') or h.get('to')}" for h in hops) # First, search primary pair's alternative shortest paths if pairs_to_search: p_a, p_b = pairs_to_search[0] primary_res = self.explain_connection(p_a, p_b, user_id) if primary_res.get("found"): p_hops = primary_res["path"] sig = get_signature(p_hops) if sig not in seen_signatures: seen_signatures.add(sig) all_unique_paths.append({ "hops": p_hops, "explanation": primary_res["explanation"] }) # Also check other shortest paths between the same pair alt_shortest = self.find_alternative_paths(p_a, p_b, user_id) for path in alt_shortest: sig = get_signature(path) if sig not in seen_signatures: seen_signatures.add(sig) all_unique_paths.append({ "hops": path, "explanation": "" # Empty explanation for on-demand generation }) # Second, search paths for other query-expanded entity pairs for p_a, p_b in pairs_to_search[1:4]: # Limit to top 3 variations try: cypher = ( "MATCH p = shortestPath(" " (a:Entity {name: $entity_a, user_id: $user_id})-[*..5]-(b:Entity {name: $entity_b, user_id: $user_id})" ") " "RETURN p" ) records = self.neo4j_client.execute_query(cypher, { "entity_a": p_a, "entity_b": p_b, "user_id": str(user_id) }) if records and records[0].get("p"): path_obj = records[0]["p"] nodes, relationships = self._get_nodes_and_rels(path_obj) path_steps = [] for i in range(len(relationships)): now_name = dict(nodes[i]).get("name", "Unknown") next_name = dict(nodes[i + 1]).get("name", "Unknown") now_type = dict(nodes[i]).get("type", "Entity") next_type = dict(nodes[i + 1]).get("type", "Entity") rel = relationships[i] rel_type = self._get_rel_type(rel) source_doc = "" page = 1 if isinstance(rel, dict): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) elif hasattr(rel, "get"): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) else: try: source_doc = dict(rel).get("source_doc", "") page = dict(rel).get("page", 1) except Exception: pass chunk_text = "" if source_doc and page: chunk_text = self.vector_retriever.get_chunk_by_page(source_doc, page, user_id) or "" path_steps.append({ "source": now_name, "source_type": now_type, "target": next_name, "target_type": next_type, "type": rel_type, "source_doc": source_doc, "page": page, "chunk_text": chunk_text }) sig = get_signature(path_steps) if sig not in seen_signatures: seen_signatures.add(sig) all_unique_paths.append({ "hops": path_steps, "explanation": "" # Empty for on-demand generation }) except Exception as e: logger.error("Failed to query path for expanded pair (%s, %s): %s", p_a, p_b, str(e)) # Third, if no paths found, run hybrid vector fallback search if not all_unique_paths and query: logger.info("No direct paths found. Executing hybrid vector-graph search fallback...") try: chunks = self.vector_retriever.retrieve_relevant_chunks(query, user_id, limit=5) # Fetch all existing entity names for this user from Neo4j all_entities_cypher = ( "MATCH (e:Entity {user_id: $user_id}) " "RETURN DISTINCT e.name AS name, e.type AS type" ) entity_records = self.neo4j_client.execute_query(all_entities_cypher, {"user_id": str(user_id)}) known_entities = {r["name"]: r.get("type", "Entity") for r in entity_records if r.get("name")} # Scan chunks to find mentioned known entities mentioned_entities = [] for chunk in chunks: chunk_text = chunk.get("text", "") for name in known_entities: if name.lower() in chunk_text.lower(): if name not in mentioned_entities: mentioned_entities.append(name) if len(mentioned_entities) >= 2: logger.info("Hybrid fallback found entities in chunks: %s", mentioned_entities) found_fallback = False for i in range(len(mentioned_entities)): for j in range(i + 1, min(len(mentioned_entities), 4)): ent_a = mentioned_entities[i] ent_b = mentioned_entities[j] cypher = ( "MATCH p = shortestPath(" " (a:Entity {name: $entity_a, user_id: $user_id})-[*..5]-(b:Entity {name: $entity_b, user_id: $user_id})" ") " "RETURN p" ) records = self.neo4j_client.execute_query(cypher, { "entity_a": ent_a, "entity_b": ent_b, "user_id": str(user_id) }) if records and records[0].get("p"): path_obj = records[0]["p"] nodes, relationships = self._get_nodes_and_rels(path_obj) path_steps = [] for k in range(len(relationships)): now_name = dict(nodes[k]).get("name", "Unknown") next_name = dict(nodes[k + 1]).get("name", "Unknown") now_type = dict(nodes[k]).get("type", "Entity") next_type = dict(nodes[k + 1]).get("type", "Entity") rel = relationships[k] rel_type = self._get_rel_type(rel) source_doc = "" page = 1 if isinstance(rel, dict): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) elif hasattr(rel, "get"): source_doc = rel.get("source_doc", "") page = rel.get("page", 1) else: try: source_doc = dict(rel).get("source_doc", "") page = dict(rel).get("page", 1) except Exception: pass chunk_txt = "" if source_doc and page: chunk_txt = self.vector_retriever.get_chunk_by_page(source_doc, page, user_id) or "" path_steps.append({ "source": now_name, "source_type": now_type, "target": next_name, "target_type": next_type, "type": rel_type, "source_doc": source_doc, "page": page, "chunk_text": chunk_txt }) sig = get_signature(path_steps) if sig not in seen_signatures: seen_signatures.add(sig) path_details = "\n".join( f"({step['source']} [{step.get('source_type', 'Entity')}]) --[{step['type']}]--> ({step['target']} [{step.get('target_type', 'Entity')}])" for step in path_steps ) try: explain_res = self.chain.invoke({ "entity_a": ent_a, "entity_b": ent_b, "path_details": path_details }) explanation = explain_res.content.strip() except Exception: explanation = f"Connection between {ent_a} and {ent_b}." all_unique_paths.append({ "hops": path_steps, "explanation": explanation }) found_fallback = True # Query alternative paths between the fallback pair alt_paths = self.find_alternative_paths(ent_a, ent_b, user_id) for alt in alt_paths: alt_sig = get_signature(alt) if alt_sig not in seen_signatures: seen_signatures.add(alt_sig) all_unique_paths.append({ "hops": alt, "explanation": "" }) break if found_fallback: break except Exception as fallback_err: logger.error("Error in hybrid fallback search: %s", str(fallback_err), exc_info=True) if not all_unique_paths: return { "found": False, "explanation": "No connection path was found in the knowledge graph.", "path": [], "alternative_paths": [], "hop_count": 0 } primary = all_unique_paths[0] alternatives = all_unique_paths[1:4] # Limit alternatives to top 3 return { "found": True, "explanation": primary["explanation"], "path": primary["hops"], "alternative_paths": alternatives, "hop_count": len(primary["hops"]) }