Coverage for src/qdrant_loader_mcp_server/fastmcp_tools/search.py: 51%
51 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""FastMCP `search` tool — semantic search across data sources."""
3from __future__ import annotations
5import asyncio
6import inspect
7from typing import Annotated, Any, Literal
9from fastmcp import Context, FastMCP
10from pydantic import BaseModel, Field
12SourceType = Literal["git", "confluence", "jira", "documentation", "localfile"]
15class HierarchyFilter(BaseModel):
16 """Hierarchy-aware filters for Confluence/localfile navigation."""
18 depth: int | None = Field(
19 default=None, description="Filter by specific hierarchy depth (0 = root pages)"
20 )
21 parent_title: str | None = Field(
22 default=None, description="Filter by parent page title"
23 )
24 root_only: bool | None = Field(
25 default=None, description="Show only root pages (no parent)"
26 )
27 has_children: bool | None = Field(
28 default=None, description="Filter by whether pages have children"
29 )
32class AttachmentFilter(BaseModel):
33 """Filters for file-attachment search."""
35 attachments_only: bool | None = Field(
36 default=None, description="Show only file attachments"
37 )
38 parent_document_title: str | None = Field(
39 default=None, description="Filter by parent document title"
40 )
41 file_type: str | None = Field(
42 default=None, description="Filter by file type (e.g., 'pdf', 'xlsx', 'png')"
43 )
44 file_size_min: int | None = Field(
45 default=None, ge=0, description="Minimum file size in bytes"
46 )
47 file_size_max: int | None = Field(
48 default=None, ge=0, description="Maximum file size in bytes"
49 )
50 author: str | None = Field(default=None, description="Filter by attachment author")
53def register_search_tools(mcp: FastMCP) -> None:
54 @mcp.tool(annotations={"readOnlyHint": True})
55 async def search(
56 ctx: Context,
57 query: Annotated[
58 str, Field(description="The search query in natural language", min_length=1)
59 ],
60 source_types: Annotated[
61 list[SourceType] | None,
62 Field(description="Optional list of source types to filter results"),
63 ] = None,
64 project_ids: Annotated[
65 list[str] | None,
66 Field(description="Optional list of project IDs to filter results"),
67 ] = None,
68 limit: Annotated[
69 int, Field(description="Maximum number of results to return", ge=1)
70 ] = 5,
71 ) -> dict[str, Any]:
72 """Perform semantic search across multiple data sources."""
73 handler = ctx.lifespan_context["search_handler"]
74 st = source_types or []
75 pids = project_ids or []
77 processed = await handler.query_processor.process_query(query)
78 results = await handler.search_engine.search(
79 query=processed["query"],
80 source_types=st,
81 project_ids=pids,
82 limit=limit,
83 )
84 if handler.reranker:
85 results = await asyncio.to_thread(
86 handler.reranker.rerank,
87 query=query,
88 results=results,
89 top_k=limit,
90 text_key="text",
91 )
93 return {
94 "results": handler.formatters.create_structured_search_results(results),
95 "total_found": len(results),
96 "query_context": {
97 "original_query": query,
98 "source_types_filtered": st,
99 "project_ids_filtered": pids,
100 },
101 }
103 @mcp.tool(annotations={"readOnlyHint": True})
104 async def hierarchy_search(
105 ctx: Context,
106 query: Annotated[
107 str, Field(description="The search query in natural language", min_length=1)
108 ],
109 hierarchy_filter: Annotated[
110 HierarchyFilter | None,
111 Field(description="Hierarchy-aware filters"),
112 ] = None,
113 organize_by_hierarchy: Annotated[
114 bool, Field(description="Group results by hierarchy structure")
115 ] = False,
116 limit: Annotated[
117 int, Field(description="Maximum number of results to return", ge=1)
118 ] = 10,
119 ) -> dict[str, Any]:
120 """Search Confluence documents with hierarchy-aware filtering and organization."""
121 handler = ctx.lifespan_context["search_handler"]
122 hf = hierarchy_filter.model_dump(exclude_none=True) if hierarchy_filter else {}
124 processed = await handler.query_processor.process_query(query)
125 results = await handler.search_engine.search(
126 query=processed["query"],
127 source_types=["confluence", "localfile"],
128 limit=max(limit * 2, 40),
129 )
130 # Filter helper may be sync or async-patched (matches legacy handling).
131 maybe = handler._apply_hierarchy_filters(results, hf)
132 filtered = await maybe if inspect.isawaitable(maybe) else maybe
133 filtered = filtered[: max(limit, 20)]
135 organized = (
136 handler._organize_by_hierarchy(filtered) if organize_by_hierarchy else None
137 )
138 return handler.formatters.create_lightweight_hierarchy_results(
139 filtered, organized or {}, query
140 )
142 @mcp.tool(annotations={"readOnlyHint": True})
143 async def attachment_search(
144 ctx: Context,
145 query: Annotated[
146 str, Field(description="The search query in natural language", min_length=1)
147 ],
148 attachment_filter: Annotated[
149 AttachmentFilter | None,
150 Field(description="Attachment filters"),
151 ] = None,
152 include_parent_context: Annotated[
153 bool, Field(description="Include parent document information in results")
154 ] = True, # TODO: Add behavior; logged-only for now
155 limit: Annotated[
156 int, Field(description="Maximum number of results to return", ge=1)
157 ] = 10,
158 ) -> dict[str, Any]:
159 """Search for file attachments and their parent documents."""
160 handler = ctx.lifespan_context["search_handler"]
161 af = (
162 attachment_filter.model_dump(exclude_none=True) if attachment_filter else {}
163 )
165 processed = await handler.query_processor.process_query(query)
166 results = await handler.search_engine.search(
167 query=processed["query"],
168 source_types=None,
169 limit=limit * 2,
170 )
171 filtered = handler._apply_lightweight_attachment_filters(results, af)
172 filtered = filtered[: max(limit, 15)]
174 attachment_groups = (
175 handler.formatters._organize_attachments_by_type(filtered)
176 if filtered
177 else []
178 )
179 return handler.formatters.create_lightweight_attachment_results(
180 attachment_groups, query
181 )