1  #!/usr/bin/env python3
  2  """
  3  Chihiro Knowledge Query Tool
  4  Efficiently load context from the knowledge database
  5  """
  6  
  7  import sqlite3
  8  import json
  9  import sys
 10  from typing import List, Dict, Tuple
 11  
 12  DB_PATH = "chihiro_knowledge.db"
 13  
 14  def load_context(context_name: str, db_path: str = DB_PATH) -> Tuple[List[Dict], int]:
 15      """Load a specific context bundle"""
 16      conn = sqlite3.connect(db_path)
 17      conn.row_factory = sqlite3.Row
 18      cursor = conn.cursor()
 19      
 20      # Get context
 21      cursor.execute("SELECT * FROM contexts WHERE name = ?", (context_name,))
 22      context = cursor.fetchone()
 23      
 24      if not context:
 25          return [], 0
 26      
 27      # Parse insight IDs
 28      insight_ids = json.loads(context['insight_ids'])
 29      
 30      # Load insights
 31      placeholders = ','.join('?' * len(insight_ids))
 32      cursor.execute(f"""
 33          SELECT id, content, source_doc, category, confidence, importance
 34          FROM insights
 35          WHERE id IN ({placeholders})
 36          ORDER BY importance DESC
 37      """, insight_ids)
 38      
 39      insights = [dict(row) for row in cursor.fetchall()]
 40      
 41      # Update access counts
 42      for insight_id in insight_ids:
 43          cursor.execute("""
 44              UPDATE insights 
 45              SET accessed_count = accessed_count + 1,
 46                  last_accessed = CURRENT_TIMESTAMP
 47              WHERE id = ?
 48          """, (insight_id,))
 49      
 50      conn.commit()
 51      conn.close()
 52      
 53      return insights, context['size_estimate']
 54  
 55  def load_all_contexts(token_budget: int = 10000, db_path: str = DB_PATH) -> Tuple[List[Dict], int]:
 56      """Load contexts by priority until token budget exhausted"""
 57      conn = sqlite3.connect(db_path)
 58      conn.row_factory = sqlite3.Row
 59      cursor = conn.cursor()
 60      
 61      # Get all contexts ordered by priority
 62      cursor.execute("SELECT * FROM contexts ORDER BY load_priority DESC")
 63      contexts = cursor.fetchall()
 64      
 65      all_insights = []
 66      tokens_used = 0
 67      contexts_loaded = []
 68      
 69      for context in contexts:
 70          if tokens_used + context['size_estimate'] <= token_budget:
 71              insights, size = load_context(context['name'], db_path)
 72              all_insights.extend(insights)
 73              tokens_used += size
 74              contexts_loaded.append(context['name'])
 75          else:
 76              break
 77      
 78      conn.close()
 79      
 80      print(f"Loaded {len(contexts_loaded)} contexts: {', '.join(contexts_loaded)}", file=sys.stderr)
 81      print(f"Total insights: {len(all_insights)}, Estimated tokens: {tokens_used}", file=sys.stderr)
 82      
 83      return all_insights, tokens_used
 84  
 85  def format_insights(insights: List[Dict]) -> str:
 86      """Format insights for context loading"""
 87      output = []
 88      
 89      # Group by category
 90      by_category = {}
 91      for insight in insights:
 92          cat = insight['category']
 93          if cat not in by_category:
 94              by_category[cat] = []
 95          by_category[cat].append(insight)
 96      
 97      # Format each category
 98      for category, items in sorted(by_category.items()):
 99          output.append(f"\n## {category.upper()}")
100          for item in items:
101              conf = "✓" if item['confidence'] >= 0.9 else "~"
102              output.append(f"{conf} {item['content']}")
103              if item['source_doc']:
104                  output.append(f"  (from {item['source_doc']})")
105      
106      return '\n'.join(output)
107  
108  def query_insights(search_term: str, limit: int = 10, db_path: str = DB_PATH) -> List[Dict]:
109      """Search insights by content"""
110      conn = sqlite3.connect(db_path)
111      conn.row_factory = sqlite3.Row
112      cursor = conn.cursor()
113      
114      cursor.execute("""
115          SELECT * FROM insights
116          WHERE content LIKE ?
117          ORDER BY importance DESC, confidence DESC
118          LIMIT ?
119      """, (f'%{search_term}%', limit))
120      
121      results = [dict(row) for row in cursor.fetchall()]
122      conn.close()
123      
124      return results
125  
126  def main():
127      if len(sys.argv) < 2:
128          print("Usage:")
129          print("  chihiro_query.py load <context_name>  - Load specific context")
130          print("  chihiro_query.py all [token_budget]   - Load all contexts by priority")
131          print("  chihiro_query.py search <term>        - Search insights")
132          sys.exit(1)
133      
134      command = sys.argv[1]
135      
136      if command == "load":
137          if len(sys.argv) < 3:
138              print("Error: context name required")
139              sys.exit(1)
140          insights, tokens = load_context(sys.argv[2])
141          print(format_insights(insights))
142          
143      elif command == "all":
144          budget = int(sys.argv[2]) if len(sys.argv) > 2 else 10000
145          insights, tokens = load_all_contexts(budget)
146          print(format_insights(insights))
147          
148      elif command == "search":
149          if len(sys.argv) < 3:
150              print("Error: search term required")
151              sys.exit(1)
152          results = query_insights(sys.argv[2])
153          print(format_insights(results))
154          
155      else:
156          print(f"Unknown command: {command}")
157          sys.exit(1)
158  
159  if __name__ == "__main__":
160      main()