Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 80 additions & 11 deletions src/unity_docs_mcp/mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import os
import sys
from pathlib import Path
from typing import List, Optional
from typing import Any, List, Optional

from mcp.server.fastmcp import FastMCP

Expand Down Expand Up @@ -76,17 +76,17 @@ def _response_meta(docstore: DocStore, baked_manifest: Optional[dict] = None) ->
return meta


@app.tool()
def search(
query: str,
k: int = 6,
source_types: Optional[List[str] | str] = None,
) -> List[dict]:
docstore = _get_docstore()
meta = _response_meta(docstore)
def _parse_source_types(source_types: Optional[List[str] | str]) -> Optional[List[str]]:
if source_types is None:
return None
if isinstance(source_types, str):
source_types = [s.strip() for s in source_types.split(",") if s.strip()]
results = docstore.search(query=query, k=k, source_types=source_types)
values = [s.strip().lower() for s in source_types.split(",") if s.strip()]
else:
values = [str(s).strip().lower() for s in source_types if str(s).strip()]
return values or None


def _serialize_search_results(results: list[Any], meta: dict) -> List[dict]:
return [
{
"chunk_id": r.chunk_id,
Expand All @@ -104,6 +104,65 @@ def search(
]


@app.tool()
def search(
query: str,
k: int = 6,
source_types: Optional[List[str] | str] = None,
debug: bool = False,
) -> List[dict] | dict:
docstore = _get_docstore()
meta = _response_meta(docstore)
parsed_source_types = _parse_source_types(source_types)
available_source_types = docstore.available_source_types()
known_source_types = docstore.known_source_types()

invalid_source_types = []
unavailable_source_types = []
if parsed_source_types:
requested_set = set(parsed_source_types)
known_set = set(known_source_types)
available_set = set(available_source_types)
invalid_source_types = sorted(requested_set - known_set)
unavailable_source_types = sorted((requested_set & known_set) - available_set)

if invalid_source_types or unavailable_source_types:
message_parts = []
if invalid_source_types:
message_parts.append(f"Unsupported source_types: {', '.join(invalid_source_types)}")
if unavailable_source_types:
message_parts.append(f"Requested source_types not present in this index: {', '.join(unavailable_source_types)}")
return {
"error": "invalid_source_types",
"message": ". ".join(message_parts),
"requested_source_types": parsed_source_types or [],
"invalid_source_types": invalid_source_types,
"unavailable_source_types": unavailable_source_types,
"known_source_types": known_source_types,
"available_source_types": available_source_types,
"results": [],
"meta": meta,
}

results = docstore.search(query=query, k=k, source_types=parsed_source_types)
serialized = _serialize_search_results(results, meta)
if not debug:
return serialized
return {
"results": serialized,
"meta": meta,
"debug": {
"query": query,
"k": k,
"requested_source_types": parsed_source_types or [],
"available_source_types": available_source_types,
"known_source_types": known_source_types,
"retrieval_mode": meta["retrieval_mode"],
"result_count": len(serialized),
},
}


@app.tool()
def open(
doc_id: Optional[str] = None,
Expand Down Expand Up @@ -177,13 +236,23 @@ def status() -> dict:
paths = make_paths(config)
baked_manifest = _read_manifest(paths.baked_dir / "manifest.json")
index_manifest = _read_manifest(paths.index_dir / "manifest.json")
available_source_types = docstore.available_source_types()
source_type_counts = docstore.source_type_counts()
coverage_warnings: list[str] = []
if "scriptref" not in set(available_source_types):
coverage_warnings.append(
"scriptref source type is not present in the loaded corpus/index; API symbol lookup coverage may be incomplete"
)
return {
"meta": _response_meta(docstore, baked_manifest=baked_manifest),
"paths": vars(config.paths),
"unity_version": config.unity_version,
"embedder": vars(config.index.embedder),
"baked_manifest": baked_manifest,
"index_manifest": index_manifest,
"available_source_types": available_source_types,
"source_type_counts": source_type_counts,
"coverage_warnings": coverage_warnings,
}


Expand Down
30 changes: 30 additions & 0 deletions src/unity_docs_mcp/tools/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
from unity_docs_mcp.index.search import HybridSearcher
from unity_docs_mcp.paths import make_paths

_DEFAULT_SOURCE_TYPES = ("manual", "scriptref")


@dataclass
class DocRecord:
Expand All @@ -26,8 +28,10 @@ def __init__(self, config: Config):
self.config = config
self.paths = make_paths(config)
self.corpus = self._load_corpus(self.paths.baked_dir / "corpus.jsonl")
self._doc_source_type_counts = self._count_source_types(self.corpus.values())
self.link_index = self._load_links(self.paths.baked_dir / "link_graph.jsonl")
self.searcher = HybridSearcher(config, self.paths.index_dir)
self._chunk_source_type_counts = self._count_source_types(getattr(self.searcher, "chunk_meta", {}).values())

def _load_corpus(self, path: Path) -> Dict[str, DocRecord]:
records: Dict[str, DocRecord] = {}
Expand Down Expand Up @@ -56,6 +60,19 @@ def _load_links(self, path: Path) -> Dict[str, List[str]]:
links.setdefault(row["from_doc_id"], []).append(row["to_doc_id"])
return links

@staticmethod
def _count_source_types(rows) -> Dict[str, int]:
counts: Dict[str, int] = {}
for row in rows:
if isinstance(row, dict):
source_type = row.get("source_type", "")
else:
source_type = getattr(row, "source_type", "")
if not source_type:
continue
counts[source_type] = counts.get(source_type, 0) + 1
return counts

def open_doc(self, doc_id: Optional[str] = None, path: Optional[str] = None) -> Optional[DocRecord]:
target_id = doc_id
if not target_id and path:
Expand All @@ -82,3 +99,16 @@ def related(self, doc_id: str, limit: int = 10) -> List[DocRecord]:

def search(self, query: str, k: int = 6, source_types: Optional[List[str]] = None) -> List:
return self.searcher.search(query=query, k=k, source_types=source_types)

def available_source_types(self) -> List[str]:
all_types = set(self._doc_source_type_counts) | set(self._chunk_source_type_counts)
return sorted(all_types)

def known_source_types(self) -> List[str]:
return sorted(set(_DEFAULT_SOURCE_TYPES) | set(self.available_source_types()))

def source_type_counts(self) -> Dict[str, Dict[str, int]]:
return {
"docs": dict(sorted(self._doc_source_type_counts.items())),
"chunks": dict(sorted(self._chunk_source_type_counts.items())),
}
39 changes: 39 additions & 0 deletions tests/test_mcp_meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
class _FakeDocStore:
def __init__(self) -> None:
self.config = Config()
self._available_source_types = ["manual"]

def search(self, query: str, k: int = 6, source_types=None):
return [
Expand Down Expand Up @@ -57,6 +58,15 @@ def related(self, doc_id: str, limit: int = 10):
)
]

def available_source_types(self):
return list(self._available_source_types)

def known_source_types(self):
return ["manual", "scriptref"]

def source_type_counts(self):
return {"docs": {"manual": 1}, "chunks": {"manual": 1}}


def _install_fake_docstore(monkeypatch):
fake = _FakeDocStore()
Expand Down Expand Up @@ -91,6 +101,32 @@ def test_open_includes_meta_even_when_missing(monkeypatch):
assert "doc_id" not in missing


def test_search_invalid_source_types_returns_actionable_error(monkeypatch):
_install_fake_docstore(monkeypatch)
result = mcp_server.search("Rigidbody.AddForce", source_types="scripting")
assert result["error"] == "invalid_source_types"
assert result["invalid_source_types"] == ["scripting"]
assert "manual" in result["available_source_types"]


def test_search_unavailable_source_types_returns_actionable_error(monkeypatch):
_install_fake_docstore(monkeypatch)
result = mcp_server.search("Rigidbody.AddForce", source_types="scriptref")
assert result["error"] == "invalid_source_types"
assert result["invalid_source_types"] == []
assert result["unavailable_source_types"] == ["scriptref"]


def test_search_debug_returns_results_and_debug_block(monkeypatch):
_install_fake_docstore(monkeypatch)
result = mcp_server.search("IJobParallelFor batch size", k=3, debug=True)
assert "results" in result
assert "debug" in result
assert result["debug"]["query"] == "IJobParallelFor batch size"
assert result["debug"]["result_count"] == 1
assert result["results"][0]["doc_id"] == "manual/job-system-parallel-for-jobs"


def test_list_files_and_related_include_meta(monkeypatch):
_install_fake_docstore(monkeypatch)
files = mcp_server.list_files("*parallel-for*")
Expand Down Expand Up @@ -118,3 +154,6 @@ def _fake_read_manifest(path):
assert status["meta"]["built_on"] == "2026-02-20"
assert "paths" in status
assert "index_manifest" in status
assert status["available_source_types"] == ["manual"]
assert status["source_type_counts"]["docs"]["manual"] == 1
assert status["coverage_warnings"]