feat(lsp): model descriptions in LSP completions (#4713)

Author: benfdking

examples/sushi/models/latest_order.sql

@@ -12,4 +12,3 @@ MODEL (
 SELECT id, customer_id, start_ts, end_ts, event_date
 FROM sushi.orders
 ORDER BY event_date DESC LIMIT 1
-

sqlmesh/lsp/completions.py

@@ -1,9 +1,14 @@
 from functools import lru_cache
 from sqlglot import Dialect, Tokenizer
-from sqlmesh.lsp.custom import AllModelsResponse, MacroCompletion
+from sqlmesh.lsp.custom import (
+    AllModelsResponse,
+    MacroCompletion,
+    ModelCompletion,
+)
 from sqlmesh import macro
 import typing as t
 from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
+from sqlmesh.lsp.description import generate_markdown_description
 from sqlmesh.lsp.uri import URI
 
 
@@ -26,38 +31,42 @@ def get_sql_completions(
     # Combine keywords - SQL keywords first, then file keywords
     all_keywords = list(sql_keywords) + list(file_keywords - sql_keywords)
 
+    models = list(get_models(context, file_uri))
     return AllModelsResponse(
-        models=list(get_models(context, file_uri)),
+        models=[m.name for m in models],
+        model_completions=models,
         keywords=all_keywords,
         macros=list(get_macros(context, file_uri)),
     )
 
 
-def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:
+def get_models(
+    context: t.Optional[LSPContext], file_uri: t.Optional[URI]
+) -> t.List[ModelCompletion]:
     """
     Return a list of models for a given file.
 
     If there is no context, return an empty list.
     If there is a context, return a list of all models bar the ones the file itself defines.
     """
     if context is None:
-        return set()
+        return []
+
+    current_path = file_uri.to_path() if file_uri is not None else None
+
+    completions: t.List[ModelCompletion] = []
+    for model in context.context.models.values():
+        if current_path is not None and model._path == current_path:
+            continue
+        description = None
+        try:
+            description = generate_markdown_description(model)
+        except Exception:
+            description = getattr(model, "description", None)
+
+        completions.append(ModelCompletion(name=model.name, description=description))
 
-    all_models = set()
-    # Extract model names from ModelInfo objects
-    for file_info in context.map.values():
-        if isinstance(file_info, ModelTarget):
-            all_models.update(file_info.names)
-
-    # Remove models from the current file
-    path = file_uri.to_path() if file_uri is not None else None
-    if path is not None and path in context.map:
-        file_info = context.map[path]
-        if isinstance(file_info, ModelTarget):
-            for model in file_info.names:
-                all_models.discard(model)
-
-    return all_models
+    return completions
 
 
 def get_macros(

sqlmesh/lsp/custom.py

@@ -30,12 +30,19 @@ class MacroCompletion(PydanticModel):
     description: t.Optional[str] = None
 
 
+class ModelCompletion(PydanticModel):
+    """Information about a model for autocompletion."""
+
+    name: str
+    description: t.Optional[str] = None
+
+
 class AllModelsResponse(CustomMethodResponseBaseClass):
-    """
-    Response to get all the models that are in the current project.
-    """
+    """Response to get all models that are in the current project."""
 
+    #: Deprecated: use ``model_completions`` instead
     models: t.List[str]
+    model_completions: t.List[ModelCompletion]
     keywords: t.List[str]
     macros: t.List[MacroCompletion]
 

sqlmesh/lsp/main.py

@@ -627,12 +627,18 @@ def completion(
 
                 completion_items = []
                 # Add model completions
-                for model in completion_response.models:
+                for model in completion_response.model_completions:
                     completion_items.append(
                         types.CompletionItem(
-                            label=model,
+                            label=model.name,
                             kind=types.CompletionItemKind.Reference,
                             detail="SQLMesh Model",
+                            documentation=types.MarkupContent(
+                                kind=types.MarkupKind.Markdown,
+                                value=model.description or "No description available",
+                            )
+                            if model.description
+                            else None,
                         )
                     )
                 # Add macro completions

tests/lsp/test_completions.py

@@ -41,6 +41,20 @@ def test_get_macros():
     assert add_one_macro.description
 
 
+def test_model_completions_include_descriptions():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    completions = LSPContext.get_completions(lsp_context, None)
+
+    model_entry = next(
+        (m for m in completions.model_completions if m.name == "sushi.customers"),
+        None,
+    )
+    assert model_entry is not None
+    assert model_entry.description
+
+
 def test_get_sql_completions_with_context_no_file_uri():
     context = Context(paths=["examples/sushi"])
     lsp_context = LSPContext(context)