Update modules/knowledge_graph.py

modules/knowledge_graph.py

@@ -290,6 +290,60 @@ class KnowledgeGraph:
         return related_concepts
 
     def search_with_depth(self, query: str, max_depth: int = 3, limit: int = 5) -> List[Tuple[str, float]]:
+        """Enhanced search with depth control and path weights"""
+        temp_query_node = f"_query_{query}"
+        
+        try:
+            # First add the query node and its relationships
+            query_relationships = self.expand_concept(query)
+            if not query_relationships:
+                print(f"No relationships found for query: {query}")
+                return []
+                
+            # Ensure query node is added before adding relationships
+            self.graph.add_node(temp_query_node, node_type='query')
+            self.add_weighted_relationships(temp_query_node, query_relationships)
+            
+            # Rest of the search logic...
+            image_scores = {}
+            image_nodes = [n for n, attr in self.graph.nodes(data=True) 
+                          if attr.get('node_type') == 'image']
+            
+            for image_id in image_nodes:
+                try:
+                    paths = nx.all_simple_paths(
+                        self.graph,
+                        source=temp_query_node,
+                        target=image_id,
+                        cutoff=max_depth
+                    )
+                    
+                    path_scores = []
+                    for path in paths:
+                        score = self._calculate_path_score(path)
+                        path_scores.append(score)
+                    
+                    if path_scores:
+                        image_scores[image_id] = max(path_scores)
+                        
+                except nx.NetworkXNoPath:
+                    continue
+            
+            # Sort and return results
+            sorted_results = sorted(
+                image_scores.items(),
+                key=lambda x: x[1],
+                reverse=True
+            )
+            return sorted_results[:limit]
+            
+        finally:
+            # Clean up: remove temporary query node
+            if self.graph.has_node(temp_query_node):
+                self.graph.remove_node(temp_query_node)
+
+    
+    def search_with_depth1(self, query: str, max_depth: int = 3, limit: int = 5) -> List[Tuple[str, float]]:
         """
         Enhanced search with depth control and path weights