Update pageindex/core/tree_index.py

pageindex/core/tree_index.py

@@ -1,226 +1,226 @@
-import json
-import re
-import asyncio
-from typing import List, Dict, Any, Optional
-
-# Import potentially available util functions or re-implement if necessary
-# Assuming we are in pageindex/core/tree_index.py, and page_index_md is in pageindex/page_index_md.py
-try:
-    from ...pageindex.page_index_md import (
-        extract_nodes_from_markdown,
-        extract_node_text_content,
-        build_tree_from_nodes,
-        count_tokens  # Assuming this exists in utils or page_index_md imports
-    )
-except ImportError:
-    # If imports fail due to path issues, we might need to adjust python path or copy logic.
-    # For now, let's assume we can import from the sibling directory if the package structure is correct.
-    # Alternatively, we can re-implement the core logic here to be self-contained.
-    pass
-
-class TreeIndex:
-    def __init__(self, model: str = "gpt-4-turbo"):
-        self.model = model
-        self.tree = None
-
-    def build_from_markdown(self, markdown_text: str) -> str:
-        """
-        Builds the tree index from markdown text.
-        
-        Args:
-            markdown_text (str): The markdown content.
-            
-        Returns:
-            str: A document ID or confirmation (for this implementation, it returns 'doc_id').
-        """
-        print(f"Extracting nodes from markdown content...")
-        # Re-implementing core logic here to ensure stability and avoid complex relative imports for now
-        # This matches page_index_md.py logic
-        
-        # 1. Extract nodes
-        node_list, lines = self._extract_nodes(markdown_text)
-        
-        # 2. Extract text content
-        nodes_with_content = self._extract_node_text(node_list, lines)
-        
-        # 3. Build tree
-        self.tree = self._build_tree(nodes_with_content)
-        
-        # 4. Add Node IDs (simple counter)
-        self._add_node_ids(self.tree)
-
-        print(f"Tree built successfully.")
-        return "doc_1" # In a real app this would be a UUID
-
-    def _extract_nodes(self, markdown_content):
-        header_pattern = r'^(#{1,6})\s+(.+)$'
-        code_block_pattern = r'^```'
-        node_list = []
-        
-        lines = markdown_content.split('\n')
-        in_code_block = False
-        
-        for line_num, line in enumerate(lines, 1):
-            stripped_line = line.strip()
-            
-            if re.match(code_block_pattern, stripped_line):
-                in_code_block = not in_code_block
-                continue
-            
-            if not stripped_line:
-                continue
-            
-            if not in_code_block:
-                match = re.match(header_pattern, stripped_line)
-                if match:
-                    title = match.group(2).strip()
-                    node_list.append({'node_title': title, 'line_num': line_num})
-        
-        return node_list, lines
-
-    def _extract_node_text(self, node_list, lines):
-        all_nodes = []
-        for node in node_list:
-            line_content = lines[node['line_num'] - 1]
-            header_match = re.match(r'^(#{1,6})', line_content)
-            
-            if header_match:
-                all_nodes.append({
-                    'title': node['node_title'],
-                    'line_num': node['line_num'],
-                    'level': len(header_match.group(1))
-                })
-        
-        for i, node in enumerate(all_nodes):
-            start_line = node['line_num'] - 1 
-            if i + 1 < len(all_nodes):
-                end_line = all_nodes[i + 1]['line_num'] - 1 
-            else:
-                end_line = len(lines)
-            
-            node['text'] = '\n'.join(lines[start_line:end_line]).strip()    
-        return all_nodes
-
-    def _build_tree(self, node_list):
-        if not node_list:
-            return []
-        
-        stack = []
-        root_nodes = []
-        
-        for node in node_list:
-            current_level = node['level']
-            
-            tree_node = {
-                'title': node['title'],
-                'text': node['text'],
-                'line_num': node['line_num'],
-                'nodes': []
-            }
-            
-            while stack and stack[-1][1] >= current_level:
-                stack.pop()
-            
-            if not stack:
-                root_nodes.append(tree_node)
-            else:
-                parent_node, parent_level = stack[-1]
-                parent_node['nodes'].append(tree_node)
-            
-            stack.append((tree_node, current_level))
-        
-        return root_nodes
-
-    def _add_node_ids(self, nodes, prefix=""):
-        for i, node in enumerate(nodes):
-            node_id = f"{prefix}{i+1}"
-            node['node_id'] = node_id
-            if node.get('nodes'):
-                self._add_node_ids(node['nodes'], prefix=f"{node_id}.")
-
-    def reasoning_search(self, query: str, llm_client: Any) -> str:
-        """
-        Performs a tree search to find relevant nodes for the query.
-        
-        Args:
-            query (str): The user query.
-            llm_client: The initialized OpenAI client.
-            
-        Returns:
-            str: Context concatenated from relevant nodes.
-        """
-        if not self.tree:
-            return "Tree not index built. Please upload a document first."
-
-        # Simplify tree for prompt (remove raw text to save tokens, keep titles and IDs)
-        tree_summary = self._get_tree_structure_summary(self.tree)
-        
-        prompt = f"""
-        You are given a query and the tree structure of a document.
-        You need to find all nodes that are likely to contain the answer.
-
-        Query: {query}
-
-        Document tree structure: 
-        {json.dumps(tree_summary, indent=2)}
-
-        Reply in the following JSON format:
-        {{
-          "thinking": <your reasoning about which nodes are relevant>,
-          "node_list": ["node_id1", "node_id2", ...]
-        }}
-        """
-
-        try:
-            response = llm_client.chat.completions.create(
-                model="gpt-4-turbo", # Provide a default, though the client might override or ignored it for NIM
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.1
-            )
-            
-            content = response.choices[0].message.content
-            # Basic JSON extraction if LLM adds markdown
-            if "```json" in content:
-                content = content.split("```json")[1].split("```")[0].strip()
-            elif "```" in content:
-                content = content.split("```")[1].split("```")[0].strip()
-            
-            result = json.loads(content)
-            relevant_node_ids = result.get("node_list", [])
-            
-            # Retrieve text for these nodes
-            context_parts = []
-            for node_id in relevant_node_ids:
-                node = self._find_node_by_id(self.tree, node_id)
-                if node:
-                    context_parts.append(f"--- Section: {node['title']} (ID: {node_id}) ---\n{node['text']}\n")
-            
-            full_context = "\n".join(context_parts)
-            if not full_context:
-                return "No relevant context found in the document tree."
-            return full_context
-
-        except Exception as e:
-            return f"Error during reasoning search: {str(e)}"
-
-    def _get_tree_structure_summary(self, nodes):
-        summary = []
-        for node in nodes:
-            node_summary = {
-                'title': node['title'],
-                'node_id': node['node_id']
-            }
-            if node.get('nodes'):
-                node_summary['nodes'] = self._get_tree_structure_summary(node['nodes'])
-            summary.append(node_summary)
-        return summary
-
-    def _find_node_by_id(self, nodes, target_id):
-        for node in nodes:
-            if node.get('node_id') == target_id:
-                return node
-            if node.get('nodes'):
-                found = self._find_node_by_id(node['nodes'], target_id)
-                if found:
-                    return found
-        return None
+import json
+import re
+import asyncio
+from typing import List, Dict, Any, Optional
+
+# Import potentially available util functions or re-implement if necessary
+# Assuming we are in pageindex/core/tree_index.py, and page_index_md is in pageindex/page_index_md.py
+try:
+    from ...pageindex.page_index_md import (
+        extract_nodes_from_markdown,
+        extract_node_text_content,
+        build_tree_from_nodes,
+        count_tokens  # Assuming this exists in utils or page_index_md imports
+    )
+except ImportError:
+    # If imports fail due to path issues, we might need to adjust python path or copy logic.
+    # For now, let's assume we can import from the sibling directory if the package structure is correct.
+    # Alternatively, we can re-implement the core logic here to be self-contained.
+    pass
+
+class TreeIndex:
+    def __init__(self, model: str = "gpt-4-turbo"):
+        self.model = model
+        self.tree = None
+
+    def build_from_markdown(self, markdown_text: str) -> str:
+        """
+        Builds the tree index from markdown text.
+        
+        Args:
+            markdown_text (str): The markdown content.
+            
+        Returns:
+            str: A document ID or confirmation (for this implementation, it returns 'doc_id').
+        """
+        print(f"Extracting nodes from markdown content...")
+        # Re-implementing core logic here to ensure stability and avoid complex relative imports for now
+        # This matches page_index_md.py logic
+        
+        # 1. Extract nodes
+        node_list, lines = self._extract_nodes(markdown_text)
+        
+        # 2. Extract text content
+        nodes_with_content = self._extract_node_text(node_list, lines)
+        
+        # 3. Build tree
+        self.tree = self._build_tree(nodes_with_content)
+        
+        # 4. Add Node IDs (simple counter)
+        self._add_node_ids(self.tree)
+
+        print(f"Tree built successfully.")
+        return "doc_1" # In a real app this would be a UUID
+
+    def _extract_nodes(self, markdown_content):
+        header_pattern = r'^(#{1,6})\s+(.+)$'
+        code_block_pattern = r'^```'
+        node_list = []
+        
+        lines = markdown_content.split('\n')
+        in_code_block = False
+        
+        for line_num, line in enumerate(lines, 1):
+            stripped_line = line.strip()
+            
+            if re.match(code_block_pattern, stripped_line):
+                in_code_block = not in_code_block
+                continue
+            
+            if not stripped_line:
+                continue
+            
+            if not in_code_block:
+                match = re.match(header_pattern, stripped_line)
+                if match:
+                    title = match.group(2).strip()
+                    node_list.append({'node_title': title, 'line_num': line_num})
+        
+        return node_list, lines
+
+    def _extract_node_text(self, node_list, lines):
+        all_nodes = []
+        for node in node_list:
+            line_content = lines[node['line_num'] - 1]
+            header_match = re.match(r'^(#{1,6})', line_content)
+            
+            if header_match:
+                all_nodes.append({
+                    'title': node['node_title'],
+                    'line_num': node['line_num'],
+                    'level': len(header_match.group(1))
+                })
+        
+        for i, node in enumerate(all_nodes):
+            start_line = node['line_num'] - 1 
+            if i + 1 < len(all_nodes):
+                end_line = all_nodes[i + 1]['line_num'] - 1 
+            else:
+                end_line = len(lines)
+            
+            node['text'] = '\n'.join(lines[start_line:end_line]).strip()    
+        return all_nodes
+
+    def _build_tree(self, node_list):
+        if not node_list:
+            return []
+        
+        stack = []
+        root_nodes = []
+        
+        for node in node_list:
+            current_level = node['level']
+            
+            tree_node = {
+                'title': node['title'],
+                'text': node['text'],
+                'line_num': node['line_num'],
+                'nodes': []
+            }
+            
+            while stack and stack[-1][1] >= current_level:
+                stack.pop()
+            
+            if not stack:
+                root_nodes.append(tree_node)
+            else:
+                parent_node, parent_level = stack[-1]
+                parent_node['nodes'].append(tree_node)
+            
+            stack.append((tree_node, current_level))
+        
+        return root_nodes
+
+    def _add_node_ids(self, nodes, prefix=""):
+        for i, node in enumerate(nodes):
+            node_id = f"{prefix}{i+1}"
+            node['node_id'] = node_id
+            if node.get('nodes'):
+                self._add_node_ids(node['nodes'], prefix=f"{node_id}.")
+
+    def reasoning_search(self, query: str, llm_client: Any, model: str = "gpt-4-turbo") -> str:
+        """
+        Performs a tree search to find relevant nodes for the query.
+        
+        Args:
+            query (str): The user query.
+            llm_client: The initialized OpenAI client.
+            
+        Returns:
+            str: Context concatenated from relevant nodes.
+        """
+        if not self.tree:
+            return "Tree not index built. Please upload a document first."
+
+        # Simplify tree for prompt (remove raw text to save tokens, keep titles and IDs)
+        tree_summary = self._get_tree_structure_summary(self.tree)
+        
+        prompt = f"""
+        You are given a query and the tree structure of a document.
+        You need to find all nodes that are likely to contain the answer.
+
+        Query: {query}
+
+        Document tree structure: 
+        {json.dumps(tree_summary, indent=2)}
+
+        Reply in the following JSON format:
+        {{
+          "thinking": <your reasoning about which nodes are relevant>,
+          "node_list": ["node_id1", "node_id2", ...]
+        }}
+        """
+
+        try:
+            response = llm_client.chat.completions.create(
+                model=model, # Use the provided model instead of hardcoded default
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.1
+            )
+            
+            content = response.choices[0].message.content
+            # Basic JSON extraction if LLM adds markdown
+            if "```json" in content:
+                content = content.split("```json")[1].split("```")[0].strip()
+            elif "```" in content:
+                content = content.split("```")[1].split("```")[0].strip()
+            
+            result = json.loads(content)
+            relevant_node_ids = result.get("node_list", [])
+            
+            # Retrieve text for these nodes
+            context_parts = []
+            for node_id in relevant_node_ids:
+                node = self._find_node_by_id(self.tree, node_id)
+                if node:
+                    context_parts.append(f"--- Section: {node['title']} (ID: {node_id}) ---\n{node['text']}\n")
+            
+            full_context = "\n".join(context_parts)
+            if not full_context:
+                return "No relevant context found in the document tree."
+            return full_context
+
+        except Exception as e:
+            return f"Error during reasoning search: {str(e)}"
+
+    def _get_tree_structure_summary(self, nodes):
+        summary = []
+        for node in nodes:
+            node_summary = {
+                'title': node['title'],
+                'node_id': node['node_id']
+            }
+            if node.get('nodes'):
+                node_summary['nodes'] = self._get_tree_structure_summary(node['nodes'])
+            summary.append(node_summary)
+        return summary
+
+    def _find_node_by_id(self, nodes, target_id):
+        for node in nodes:
+            if node.get('node_id') == target_id:
+                return node
+            if node.get('nodes'):
+                found = self._find_node_by_id(node['nodes'], target_id)
+                if found:
+                    return found
+        return None