feat: enhance CLI output and improve knowledge store handling in MCP server

Author: deepakdgupta1

.vscode/settings.json

@@ -0,0 +1 @@
+{}

src/knowcode/cli/cli.py

@@ -67,10 +67,14 @@ def analyze(directory: str, output: str, ignore: tuple[str, ...], temporal: bool
     click.echo(f"  Relationships: {stats['total_relationships']}")
     if stats.get('total_errors', 0) > 0:
         click.echo(f"  Errors: {stats['total_errors']}")
+    if stats.get("indexed_chunks") is not None:
+        click.echo(f"  Indexed chunks: {stats['indexed_chunks']}")
 
     output_path = Path(output)
     save_path = output_path / KnowledgeStore.DEFAULT_FILENAME if output_path.is_dir() else output_path
     click.echo(f"\n  Saved to: {save_path}")
+    if stats.get("index_path"):
+        click.echo(f"  Index saved to: {stats['index_path']}")
 
 
 @cli.command()
@@ -557,15 +561,12 @@ def mcp_server(store: str, config: Optional[str]) -> None:
         }
     """
     store_path = Path(store)
-    if store_path.is_dir():
-        store_path = store_path / KnowledgeStore.DEFAULT_FILENAME
-        
-    if not store_path.exists():
+    store_file = store_path / KnowledgeStore.DEFAULT_FILENAME if store_path.is_dir() else store_path
+    if not store_file.exists():
         click.echo(
-            "Error: Knowledge store not found. Run 'knowcode analyze' first.",
-            err=True
+            "⚠️ Knowledge store not found. It will be built on first query.",
+            err=True,
         )
-        sys.exit(1)
 
     try:
         from knowcode.mcp.server import run_server

src/knowcode/mcp/server.py

@@ -156,14 +156,14 @@ def __init__(self, store_path: str | Path, config_path: Optional[str] = None) ->
         self.config_path = config_path
         self._service: Optional[KnowCodeService] = None
 
-    def _ensure_service(self) -> KnowCodeService:
+    def _ensure_service(self, allow_missing_store: bool = False) -> KnowCodeService:
         """Create the shared service (loads store/index lazily)."""
         if self._service is None:
             store_file = self.store_path
             if store_file.is_dir():
                 store_file = store_file / KnowledgeStore.DEFAULT_FILENAME
 
-            if not store_file.exists():
+            if not allow_missing_store and not store_file.exists():
                 raise FileNotFoundError(
                     f"Knowledge store not found: {store_file}\n"
                     "Run 'knowcode analyze' first to build the knowledge graph."
@@ -174,6 +174,14 @@ def _ensure_service(self) -> KnowCodeService:
                 config_path=self.config_path,
             )
         return self._service
+
+    def _ensure_store_ready(self, service: KnowCodeService) -> None:
+        """Ensure the knowledge store exists by running analyze if needed."""
+        store_root = service.store_path if service.store_path.is_dir() else service.store_path.parent
+        store_file = store_root / KnowledgeStore.DEFAULT_FILENAME
+        if store_file.exists():
+            return
+        service.analyze(directory=store_root, output=store_root)
 
     def search_codebase(self, query: str, limit: int = 10) -> list[dict[str, Any]]:
         """Search for entities by name pattern.
@@ -185,7 +193,8 @@ def search_codebase(self, query: str, limit: int = 10) -> list[dict[str, Any]]:
         Returns:
             List of matching entity summaries.
         """
-        service = self._ensure_service()
+        service = self._ensure_service(allow_missing_store=True)
+        self._ensure_store_ready(service)
         entities = service.store.search(query)[:limit]
 
         return [
@@ -216,7 +225,8 @@ def get_entity_context(
         Returns:
             Context bundle with sufficiency score.
         """
-        service = self._ensure_service()
+        service = self._ensure_service(allow_missing_store=True)
+        self._ensure_store_ready(service)
         try:
             task = TaskType(task_type)
         except ValueError:
@@ -260,7 +270,8 @@ def trace_calls(
         Returns:
             List of entities with call_depth.
         """
-        service = self._ensure_service()
+        service = self._ensure_service(allow_missing_store=True)
+        self._ensure_store_ready(service)
         return service.store.trace_calls(
             entity_id,
             direction=direction,
@@ -277,7 +288,7 @@ def retrieve_context_for_query(
         expand_deps: bool = True,
     ) -> dict[str, Any]:
         """Retrieve a task-aware context bundle for a query."""
-        service = self._ensure_service()
+        service = self._ensure_service(allow_missing_store=True)
         task_override: Optional[TaskType] = None
         if task_type != "auto":
             try:

src/knowcode/service.py

@@ -125,6 +125,34 @@ def retrieve_context_for_query(
         from knowcode.llm.query_classifier import classify_query
 
         errors: list[str] = []
+        store_root = self.store_path if self.store_path.is_dir() else self.store_path.parent
+        store_file = store_root / KnowledgeStore.DEFAULT_FILENAME
+        index_path = store_root / "knowcode_index"
+
+        if not store_file.exists():
+            try:
+                self.analyze(directory=store_root, output=store_root)
+            except Exception as e:
+                return {
+                    "query": query,
+                    "task_type": task_type.value if task_type else "general",
+                    "task_confidence": 0.0,
+                    "retrieval_mode": "none",
+                    "context_text": "",
+                    "total_tokens": 0,
+                    "max_tokens": max_tokens,
+                    "truncated": False,
+                    "sufficiency_score": 0.0,
+                    "selected_entities": [],
+                    "evidence": [],
+                    "errors": [f"Auto-analyze failed: {e}"],
+                }
+
+        if not index_path.exists():
+            try:
+                self._build_index(store_root, index_path)
+            except Exception as e:
+                errors.append(f"Auto-index failed; falling back to lexical: {e}")
 
         detected_task_type, confidence = classify_query(query)
         resolved_task_type = task_type or detected_task_type
@@ -149,9 +177,6 @@ def retrieve_context_for_query(
         if per_entity_max_tokens is None:
             per_entity_max_tokens = max(200, min(2000, max_tokens // limit_entities))
 
-        store_root = self.store_path if self.store_path.is_dir() else self.store_path.parent
-        index_path = store_root / "knowcode_index"
-
         selected_entity_ids: list[str] = []
         evidence: list[dict[str, Any]] = []
         retrieval_mode = "lexical"
@@ -272,6 +297,18 @@ def add_entity_ids(items: list[dict[str, Any]]) -> None:
             "errors": errors,
         }
 
+    def _build_index(self, directory: str | Path, index_path: str | Path) -> int:
+        """Build a semantic index for a directory and persist it."""
+        from knowcode.llm.embedding import create_embedding_provider
+        from knowcode.indexing.indexer import Indexer
+
+        provider = create_embedding_provider(app_config=self.app_config)
+        indexer = Indexer(provider)
+        count = indexer.index_directory(directory)
+        indexer.save(index_path)
+        self._indexer = indexer
+        return count
+
     def _extract_query_keywords(self, query: str) -> list[str]:
         """Extract identifier-like keywords from a natural-language query."""
         stopwords = {
@@ -381,8 +418,15 @@ def analyze(
         output_path = Path(output)
         store.save(output_path)
         self._store = store
-        
-        return builder.stats()
+
+        store_root = output_path if output_path.is_dir() else output_path.parent
+        index_path = store_root / "knowcode_index"
+        index_count = self._build_index(Path(directory), index_path)
+
+        stats = builder.stats()
+        stats["indexed_chunks"] = index_count
+        stats["index_path"] = str(index_path)
+        return stats
 
     def search(self, pattern: str) -> list[dict[str, Any]]:
         """Search entities by pattern.