lean-explore 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

lean_explore/api/client.py

@@ -7,7 +7,8 @@ communication with the backend Lean Explore search engine API for
 performing searches and retrieving detailed information.
 """
 
-from typing import List, Optional
+import asyncio
+from typing import List, Optional, Union, overload
 
 import httpx
 
@@ -44,81 +45,172 @@ class Client:
         self.timeout: float = timeout
         self._headers: dict = {"Authorization": f"Bearer {self.api_key}"}
 
+    async def _fetch_one_search(
+        self,
+        client: httpx.AsyncClient,
+        query: str,
+        package_filters: Optional[List[str]],
+    ) -> APISearchResponse:
+        """Coroutine to fetch a single search result.
+
+        Args:
+            client: An active httpx.AsyncClient instance.
+            query: The search query string.
+            package_filters: An optional list of package names.
+
+        Returns:
+            An APISearchResponse object.
+        """
+        endpoint = f"{self.base_url}/search"
+        params = {"q": query}
+        if package_filters:
+            params["pkg"] = package_filters
+
+        response = await client.get(endpoint, params=params, headers=self._headers)
+        response.raise_for_status()
+        return APISearchResponse(**response.json())
+
+    @overload
     async def search(
         self, query: str, package_filters: Optional[List[str]] = None
-    ) -> APISearchResponse:
+    ) -> APISearchResponse: ...
+
+    @overload
+    async def search(
+        self, query: List[str], package_filters: Optional[List[str]] = None
+    ) -> List[APISearchResponse]: ...
+
+    async def search(
+        self,
+        query: Union[str, List[str]],
+        package_filters: Optional[List[str]] = None,
+    ) -> Union[APISearchResponse, List[APISearchResponse]]:
         """Performs a search for statement groups via the API.
 
+        This method can handle a single query string or a list of query strings.
+        When a list is provided, requests are sent concurrently.
+
         Args:
-            query: The search query string.
+            query: The search query string or a list of query strings.
             package_filters: An optional list of package names to filter the
-                search by.
+                search by. This filter is applied to all queries.
 
         Returns:
-            An APISearchResponse object containing the search results and
-            associated metadata.
+            An APISearchResponse object if a single query was provided, or a
+            list of APISearchResponse objects if a list of queries was provided.
 
         Raises:
             httpx.HTTPStatusError: If the API returns an HTTP error status (4xx or 5xx).
             httpx.RequestError: For network-related issues or other request errors.
         """
-        endpoint = f"{self.base_url}/search"
-        params = {"q": query}
-        if package_filters:
-            params["pkg"] = package_filters
+        was_single_query = isinstance(query, str)
+        queries = [query] if was_single_query else query
 
         async with httpx.AsyncClient(timeout=self.timeout) as client:
-            response = await client.get(endpoint, params=params, headers=self._headers)
-            response.raise_for_status()
-            return APISearchResponse(**response.json())
-
-    async def get_by_id(self, group_id: int) -> Optional[APISearchResultItem]:
+            tasks = [
+                self._fetch_one_search(client, q, package_filters) for q in queries
+            ]
+            results = await asyncio.gather(*tasks)
+
+        if was_single_query:
+            return results[0]
+        return results
+
+    async def _fetch_one_by_id(
+        self, client: httpx.AsyncClient, group_id: int
+    ) -> Optional[APISearchResultItem]:
+        endpoint = f"{self.base_url}/statement_groups/{group_id}"
+        response = await client.get(endpoint, headers=self._headers)
+        if response.status_code == 404:
+            return None
+        response.raise_for_status()
+        return APISearchResultItem(**response.json())
+
+    @overload
+    async def get_by_id(self, group_id: int) -> Optional[APISearchResultItem]: ...
+
+    @overload
+    async def get_by_id(
+        self, group_id: List[int]
+    ) -> List[Optional[APISearchResultItem]]: ...
+
+    async def get_by_id(
+        self, group_id: Union[int, List[int]]
+    ) -> Union[Optional[APISearchResultItem], List[Optional[APISearchResultItem]]]:
         """Retrieves a specific statement group by its unique ID via the API.
 
         Args:
-            group_id: The unique identifier of the statement group.
+            group_id: The unique identifier of the statement group, or a list of IDs.
 
         Returns:
-            An APISearchResultItem object if the statement group is found,
-            otherwise None if a 404 error is received.
+            An APISearchResultItem object if a single ID was found, None if it was
+            not found. A list of Optional[APISearchResultItem] if a list of
+            IDs was provided.
 
         Raises:
             httpx.HTTPStatusError: If the API returns an HTTP error status
-                                   other than 404 (e.g., 401, 403, 5xx).
+                                other than 404 (e.g., 401, 403, 5xx).
             httpx.RequestError: For network-related issues or other request errors.
         """
-        endpoint = f"{self.base_url}/statement_groups/{group_id}"
+        was_single_id = isinstance(group_id, int)
+        group_ids = [group_id] if was_single_id else group_id
+
         async with httpx.AsyncClient(timeout=self.timeout) as client:
-            response = await client.get(endpoint, headers=self._headers)
-            if response.status_code == 404:
-                return None
-            response.raise_for_status()
-            return APISearchResultItem(**response.json())
+            tasks = [self._fetch_one_by_id(client, g_id) for g_id in group_ids]
+            results = await asyncio.gather(*tasks)
+
+        if was_single_id:
+            return results[0]
+        return results
 
-    async def get_dependencies(self, group_id: int) -> Optional[APICitationsResponse]:
+    async def _fetch_one_dependencies(
+        self, client: httpx.AsyncClient, group_id: int
+    ) -> Optional[APICitationsResponse]:
+        endpoint = f"{self.base_url}/statement_groups/{group_id}/dependencies"
+        response = await client.get(endpoint, headers=self._headers)
+        if response.status_code == 404:
+            return None
+        response.raise_for_status()
+        return APICitationsResponse(**response.json())
+
+    @overload
+    async def get_dependencies(
+        self, group_id: int
+    ) -> Optional[APICitationsResponse]: ...
+
+    @overload
+    async def get_dependencies(
+        self, group_id: List[int]
+    ) -> List[Optional[APICitationsResponse]]: ...
+
+    async def get_dependencies(
+        self, group_id: Union[int, List[int]]
+    ) -> Union[Optional[APICitationsResponse], List[Optional[APICitationsResponse]]]:
         """Retrieves the dependencies (citations) for a specific statement group.
 
-        This method fetches the statement groups that the specified 'group_id'
-        depends on (i.e., cites).
+        This method fetches the statement groups that the specified 'group_id'(s)
+        depend on (i.e., cite).
 
         Args:
-            group_id: The unique identifier of the statement group for which
-                to fetch dependencies.
+            group_id: The unique identifier of the statement group, or a list of IDs.
 
         Returns:
-            An APICitationsResponse object containing the list of dependencies
-            (cited items) if the source statement group is found. Returns None
-            if the source statement group itself is not found (receives a 404).
+            An APICitationsResponse object if a single ID was provided. A list
+            of Optional[APICitationsResponse] if a list of IDs was provided.
+            None is returned for IDs that are not found.
 
         Raises:
             httpx.HTTPStatusError: If the API returns an HTTP error status
-                                   other than 404 (e.g., 401, 403, 5xx).
+                                other than 404 (e.g., 401, 403, 5xx).
             httpx.RequestError: For network-related issues or other request errors.
         """
-        endpoint = f"{self.base_url}/statement_groups/{group_id}/dependencies"
+        was_single_id = isinstance(group_id, int)
+        group_ids = [group_id] if was_single_id else group_id
+
         async with httpx.AsyncClient(timeout=self.timeout) as client:
-            response = await client.get(endpoint, headers=self._headers)
-            if response.status_code == 404:
-                return None
-            response.raise_for_status()
-            return APICitationsResponse(**response.json())
+            tasks = [self._fetch_one_dependencies(client, g_id) for g_id in group_ids]
+            results = await asyncio.gather(*tasks)
+
+        if was_single_id:
+            return results[0]
+        return results

lean_explore/cli/agent.py

@@ -520,19 +520,27 @@ async def _run_agent_session(
                     "to the user.\n"
                     "**Output:** CLI-friendly (plain text, simple lists). "
                     "NO complex Markdown/LaTeX.\n\n"
+                    "**Tool Usage & Efficiency:**\n"
+                    "* The `search`, `get_by_id`, and `get_dependencies` tools can "
+                    "all accept a list of inputs (queries or integer IDs) to "
+                    "perform batch operations. This is highly efficient. For "
+                    "example, `search(query=['query 1', 'query 2'])` or "
+                    "`get_by_id(group_id=[123, 456])`.\n"
+                    "* Always prefer making one batch call over multiple single "
+                    "calls.\n\n"
                     "**Packages:** Use exact top-level names for filters (Batteries, "
                     "Init, Lean, Mathlib, PhysLean, Std). Map subpackage mentions "
                     "to top-level (e.g., 'Mathlib.Analysis' -> 'Mathlib').\n\n"
                     "**Core Workflow:**\n"
                     "1.  **Search & Analyze:**\n"
                     "    * Execute multiple distinct `search` queries for each user "
-                    "request (e.g., using full statements, rephrasing). Set `limit` "
+                    "request by passing a list of queries to the tool. Set `limit` "
                     ">= 10 for each search.\n"
                     "    * From all search results, select the statement(s) most "
                     "helpful to the user.\n"
-                    "    * For each selected statement, **MUST** use "
-                    "`get_dependencies` to understand its context before "
-                    "explaining.\n\n"
+                    "    * For each selected statement, use `get_dependencies` to "
+                    "understand its context. Do this efficiently by collecting all "
+                    "relevant IDs and passing them in a single list call.\n\n"
                     "2.  **Explain Results (Conversational & CLI-Friendly):**\n"
                     "    * Briefly state your search approach (e.g., 'I looked into X "
                     "in Mathlib...').\n"
@@ -543,15 +551,14 @@ async def _run_agent_session(
                     "`informal_description`, `statement_text`).\n"
                     "        * Provide the full Lean code (`statement_text`).\n"
                     "        * Explain key dependencies (what they are, their role, "
-                    "using `statement_text` or `display_statement_text` from "
-                    "`get_dependencies` output).\n"
+                    "using `statement_text` from `get_dependencies` output).\n"
                     "3.  **Specific User Follow-ups (If Asked):**\n"
-                    "    * **`get_by_id`:** For a specific ID, provide: ID, Lean name, "
-                    "statement text, source/line, docstring, informal description "
-                    "(structured CLI format).\n"
-                    "    * **`get_dependencies` (Direct Request):** For all "
-                    "dependencies of an ID, list: ID, Lean name, statement "
-                    "text/summary. State total count.\n\n"
+                    "    * **`get_by_id`:** For one or more IDs, provide: ID, "
+                    "Lean name, statement text, source/line, docstring, informal "
+                    "description (structured CLI format).\n"
+                    "    * **`get_dependencies` (Direct Request):** For one or more "
+                    "IDs, list dependencies for each: ID, Lean name, statement "
+                    "text/summary. State total count per ID.\n\n"
                     "Always be concise, helpful, and clear."
                 ),
                 mcp_servers=[server_instance],

lean_explore/local/service.py

@@ -9,7 +9,7 @@ data assets (SQLite database, FAISS index, and embedding models).
 
 import logging
 import time
-from typing import List, Optional
+from typing import List, Optional, Union, overload
 
 import faiss  # For type hinting if needed
 from sentence_transformers import SentenceTransformer  # For type hinting if needed
@@ -69,10 +69,10 @@ class Service:
 
         Raises:
             FileNotFoundError: If essential data files (DB, FAISS index, map)
-                               are not found at their expected locations.
+                                are not found at their expected locations.
             RuntimeError: If the embedding model fails to load or if other
-                          critical initialization steps (like database connection
-                          after file checks) fail.
+                        critical initialization steps (like database connection
+                        after file checks) fail.
         """
         logger.info("Initializing local Service...")
         try:
@@ -210,43 +210,25 @@ class Service:
             informal_description=sg_orm.informal_description,
         )
 
-    def search(
+    def _perform_one_search(
         self,
         query: str,
-        package_filters: Optional[List[str]] = None,
-        limit: Optional[int] = None,
+        package_filters: Optional[List[str]],
+        limit: Optional[int],
     ) -> APISearchResponse:
-        """Performs a local search for statement groups.
+        """Helper to perform and package a single local search.
 
         Args:
             query: The search query string.
             package_filters: An optional list of package names to filter results by.
             limit: An optional limit on the number of results to return.
-                   If None, defaults.DEFAULT_RESULTS_LIMIT is used.
 
         Returns:
-            An APISearchResponse object containing search results and metadata.
-
-        Raises:
-            RuntimeError: If service not properly initialized (e.g., assets missing).
-            Exception: Propagates exceptions from `perform_search`.
+            An APISearchResponse for the given query.
         """
         start_time = time.time()
         actual_limit = limit if limit is not None else self.default_results_limit
 
-        if (
-            self.embedding_model is None
-            or self.faiss_index is None
-            or self.text_chunk_id_map is None
-        ):
-            logger.error(
-                "Search service assets not loaded. Service may not have initialized "
-                "correctly."
-            )
-            raise RuntimeError(
-                "Search service assets not loaded. Please ensure data has been fetched."
-            )
-
         with self.SessionLocal() as session:
             try:
                 ranked_results_orm = perform_search(
@@ -291,102 +273,207 @@ class Service:
             processing_time_ms=processing_time_ms,
         )
 
-    def get_by_id(self, group_id: int) -> Optional[APISearchResultItem]:
+    @overload
+    def search(
+        self,
+        query: str,
+        package_filters: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+    ) -> APISearchResponse: ...
+
+    @overload
+    def search(
+        self,
+        query: List[str],
+        package_filters: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+    ) -> List[APISearchResponse]: ...
+
+    def search(
+        self,
+        query: Union[str, List[str]],
+        package_filters: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+    ) -> Union[APISearchResponse, List[APISearchResponse]]:
+        """Performs a local search for statement groups.
+
+        This method can handle a single query string or a list of query strings.
+        When a list is provided, searches are performed serially.
+
+        Args:
+            query: The search query string or a list of query strings.
+            package_filters: An optional list of package names to filter results by.
+            limit: An optional limit on the number of results to return.
+                    If None, defaults.DEFAULT_RESULTS_LIMIT is used.
+
+        Returns:
+            An APISearchResponse object if a single query was provided, or a
+            list of APISearchResponse objects if a list of queries was provided.
+
+        Raises:
+            RuntimeError: If service not properly initialized (e.g., assets missing).
+            Exception: Propagates exceptions from `perform_search`.
+        """
+        if (
+            self.embedding_model is None
+            or self.faiss_index is None
+            or self.text_chunk_id_map is None
+        ):
+            logger.error(
+                "Search service assets not loaded. Service may not have initialized "
+                "correctly."
+            )
+            raise RuntimeError(
+                "Search service assets not loaded. Please ensure data has been fetched."
+            )
+
+        was_single_query = isinstance(query, str)
+        queries = [query] if was_single_query else query
+        results = []
+
+        for q in queries:
+            response = self._perform_one_search(q, package_filters, limit)
+            results.append(response)
+
+        if was_single_query:
+            return results[0]
+        return results
+
+    @overload
+    def get_by_id(self, group_id: int) -> Optional[APISearchResultItem]: ...
+
+    @overload
+    def get_by_id(self, group_id: List[int]) -> List[Optional[APISearchResultItem]]: ...
+
+    def get_by_id(
+        self, group_id: Union[int, List[int]]
+    ) -> Union[Optional[APISearchResultItem], List[Optional[APISearchResultItem]]]:
         """Retrieves a specific statement group by its ID from local data.
 
         Args:
-            group_id: The unique identifier of the statement group.
+            group_id: The unique identifier of the statement group, or a list of IDs.
 
         Returns:
-            An APISearchResultItem if found, otherwise None.
+            An APISearchResultItem if a single ID was found, None if not found.
+            A list of Optional[APISearchResultItem] if a list of IDs was provided.
         """
+        was_single_id = isinstance(group_id, int)
+        group_ids = [group_id] if was_single_id else group_id
+        results = []
+
         with self.SessionLocal() as session:
-            try:
-                stmt_group_orm = (
-                    session.query(StatementGroup)
-                    .options(joinedload(StatementGroup.primary_declaration))
-                    .filter(StatementGroup.id == group_id)
-                    .first()
-                )
-                if stmt_group_orm:
-                    return self._serialize_sg_to_api_item(stmt_group_orm)
-                return None
-            except SQLAlchemyError as e:
-                logger.error(
-                    f"Database error in get_by_id for group_id {group_id}: {e}",
-                    exc_info=True,
-                )
-                return None
-            except Exception as e:
-                logger.error(
-                    f"Unexpected error in get_by_id for group_id {group_id}: {e}",
-                    exc_info=True,
-                )
-                return None
+            for g_id in group_ids:
+                try:
+                    stmt_group_orm = (
+                        session.query(StatementGroup)
+                        .options(joinedload(StatementGroup.primary_declaration))
+                        .filter(StatementGroup.id == g_id)
+                        .first()
+                    )
+                    if stmt_group_orm:
+                        results.append(self._serialize_sg_to_api_item(stmt_group_orm))
+                    else:
+                        results.append(None)
+                except SQLAlchemyError as e:
+                    logger.error(
+                        f"Database error in get_by_id for group_id {g_id}: {e}",
+                        exc_info=True,
+                    )
+                    results.append(None)
+                except Exception as e:
+                    logger.error(
+                        f"Unexpected error in get_by_id for group_id {g_id}: {e}",
+                        exc_info=True,
+                    )
+                    results.append(None)
 
-    def get_dependencies(self, group_id: int) -> Optional[APICitationsResponse]:
+        if was_single_id:
+            return results[0]
+        return results
+
+    @overload
+    def get_dependencies(self, group_id: int) -> Optional[APICitationsResponse]: ...
+
+    @overload
+    def get_dependencies(
+        self, group_id: List[int]
+    ) -> List[Optional[APICitationsResponse]]: ...
+
+    def get_dependencies(
+        self, group_id: Union[int, List[int]]
+    ) -> Union[Optional[APICitationsResponse], List[Optional[APICitationsResponse]]]:
         """Retrieves citations for a specific statement group from local data.
 
         Citations are the statement groups that the specified group_id depends on.
 
         Args:
-            group_id: The unique identifier of the statement group for which
-                      to fetch citations.
+            group_id: The unique ID of the source group, or a list of IDs.
 
         Returns:
-            An APICitationsResponse object if the source group is found and has
-            citations, or an APICitationsResponse with an empty list if no
-            citations, otherwise None if the source group itself is not found or
-            a DB error occurs.
+            An APICitationsResponse if a single ID was provided, or a list of
+            Optional[APICitationsResponse] if a list of IDs was given. Returns
+            None for IDs that are not found or cause an error.
         """
+        was_single_id = isinstance(group_id, int)
+        group_ids = [group_id] if was_single_id else group_id
+        results = []
+
         with self.SessionLocal() as session:
-            try:
-                source_group_exists = (
-                    session.query(StatementGroup.id)
-                    .filter(StatementGroup.id == group_id)
-                    .first()
-                )
-                if not source_group_exists:
-                    logger.warning(
-                        f"Source statement group ID {group_id} not found for "
-                        "dependency lookup."
+            for g_id in group_ids:
+                try:
+                    source_group_exists = (
+                        session.query(StatementGroup.id)
+                        .filter(StatementGroup.id == g_id)
+                        .first()
                     )
-                    return None
-
-                cited_target_groups_orm = (
-                    session.query(StatementGroup)
-                    .join(
-                        StatementGroupDependency,
-                        StatementGroup.id
-                        == StatementGroupDependency.target_statement_group_id,
+                    if not source_group_exists:
+                        logger.warning(
+                            f"Source statement group ID {g_id} not found for "
+                            "dependency lookup."
+                        )
+                        results.append(None)
+                        continue
+
+                    cited_target_groups_orm = (
+                        session.query(StatementGroup)
+                        .join(
+                            StatementGroupDependency,
+                            StatementGroup.id
+                            == StatementGroupDependency.target_statement_group_id,
+                        )
+                        .filter(
+                            StatementGroupDependency.source_statement_group_id == g_id
+                        )
+                        .options(joinedload(StatementGroup.primary_declaration))
+                        .all()
                     )
-                    .filter(
-                        StatementGroupDependency.source_statement_group_id == group_id
-                    )
-                    .options(joinedload(StatementGroup.primary_declaration))
-                    .all()
-                )
 
-                citations_api_items = [
-                    self._serialize_sg_to_api_item(sg_orm)
-                    for sg_orm in cited_target_groups_orm
-                ]
+                    citations_api_items = [
+                        self._serialize_sg_to_api_item(sg_orm)
+                        for sg_orm in cited_target_groups_orm
+                    ]
+
+                    results.append(
+                        APICitationsResponse(
+                            source_group_id=g_id,
+                            citations=citations_api_items,
+                            count=len(citations_api_items),
+                        )
+                    )
+                except SQLAlchemyError as e:
+                    logger.error(
+                        f"Database error in get_dependencies for group_id {g_id}: {e}",
+                        exc_info=True,
+                    )
+                    results.append(None)
+                except Exception as e:
+                    logger.error(
+                        f"Unexpected error in get_dependencies for "
+                        f"group_id {g_id}: {e}",
+                        exc_info=True,
+                    )
+                    results.append(None)
 
-                return APICitationsResponse(
-                    source_group_id=group_id,
-                    citations=citations_api_items,
-                    count=len(citations_api_items),
-                )
-            except SQLAlchemyError as e:
-                logger.error(
-                    f"Database error in get_dependencies for group_id {group_id}: {e}",
-                    exc_info=True,
-                )
-                return None
-            except Exception as e:
-                logger.error(
-                    f"Unexpected error in get_dependencies for "
-                    f"group_id {group_id}: {e}",
-                    exc_info=True,
-                )
-                return None
+        if was_single_id:
+            return results[0]
+        return results

lean_explore/mcp/tools.py

@@ -10,7 +10,7 @@ made available through the MCP application context.
 
 import asyncio  # Needed for asyncio.iscoroutinefunction
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from mcp.server.fastmcp import Context as MCPContext
 
@@ -82,161 +82,189 @@ def _prepare_mcp_result_item(backend_item: APISearchResultItem) -> APISearchResu
 @mcp_app.tool()
 async def search(
     ctx: MCPContext,
-    query: str,
+    query: Union[str, List[str]],
     package_filters: Optional[List[str]] = None,
     limit: int = 10,
-) -> Dict[str, Any]:
-    """Searches Lean statement groups by a query string.
+) -> List[Dict[str, Any]]:
+    """Searches Lean statement groups by a query string or list of strings.
 
     This tool allows for filtering by package names and limits the number
-    of results returned.
+    of results returned per query.
 
     Args:
         ctx: The MCP context, providing access to shared resources like the
              backend service.
-        query: The search query string. For example, "continuous function" or
-               "prime number theorem".
+        query: A single search query string or a list of query strings. For
+               example, "continuous function" or ["prime number theorem",
+               "fundamental theorem of arithmetic"].
         package_filters: An optional list of package names to filter the search
                          results by. For example, `["Mathlib.Analysis",
                          "Mathlib.Order"]`. If None or empty, no package filter
                          is applied.
-        limit: The maximum number of search results to return from this tool.
+        limit: The maximum number of search results to return per query.
                Defaults to 10. Must be a positive integer.
 
     Returns:
-        A dictionary corresponding to the APISearchResponse model, containing
-        the search results (potentially truncated by the `limit` parameter of
-        this tool), and metadata about the search operation. The
-        `display_statement_text` field within each result item is omitted.
+        A list of dictionaries, where each dictionary corresponds to the
+        APISearchResponse model. Each response contains the search results
+        for a single query. The `display_statement_text` field within each
+        result item is omitted.
     """
     backend = await _get_backend_from_context(ctx)
     logger.info(
-        f"MCP Tool 'search' called with query: '{query}', "
+        f"MCP Tool 'search' called with query/queries: '{query}', "
         f"packages: {package_filters}, tool_limit: {limit}"
     )
 
     if not hasattr(backend, "search"):
         logger.error("Backend service does not have a 'search' method.")
-        # This should ideally return a structured error for MCP if possible.
-        # For now, FastMCP will convert this RuntimeError.
         raise RuntimeError("Search functionality not available on configured backend.")
 
-    tool_limit = max(1, limit)  # Ensure limit is at least 1 for slicing
-    api_response_pydantic: Optional[APISearchResponse]
+    tool_limit = max(1, limit)
+    backend_responses: Union[APISearchResponse, List[APISearchResponse]]
 
     # Conditionally await based on the backend's search method type
     if asyncio.iscoroutinefunction(backend.search):
-        api_response_pydantic = await backend.search(
-            query=query,
-            package_filters=package_filters,
-            # The backend.search method uses its own internal default for limit
-            # if None is passed, or the passed limit.
-            # The MCP tool will truncate the results later using tool_limit.
-        )
-    else:
-        api_response_pydantic = backend.search(
+        backend_responses = await backend.search(
             query=query, package_filters=package_filters
         )
+    else:
+        backend_responses = backend.search(query=query, package_filters=package_filters)
 
-    if not api_response_pydantic:
-        logger.warning("Backend search returned None, responding with empty results.")
-        empty_response = APISearchResponse(
-            query=query,
-            packages_applied=package_filters or [],
-            results=[],
-            count=0,
-            total_candidates_considered=0,
-            processing_time_ms=0,
+    # Normalize to a list for consistent processing, handling None from backend.
+    if backend_responses is None:
+        responses_list = []
+    else:
+        responses_list = (
+            [backend_responses]
+            if isinstance(backend_responses, APISearchResponse)
+            else backend_responses
         )
-        return empty_response.model_dump(exclude_none=True)
 
-    actual_backend_results = api_response_pydantic.results
-
-    mcp_results_list = []
-    for backend_item in actual_backend_results[:tool_limit]:  # Apply MCP tool's limit
-        mcp_results_list.append(_prepare_mcp_result_item(backend_item))
-
-    final_mcp_response = APISearchResponse(
-        query=api_response_pydantic.query,
-        packages_applied=api_response_pydantic.packages_applied,
-        results=mcp_results_list,
-        count=len(mcp_results_list),  # Count is after this tool's truncation
-        total_candidates_considered=api_response_pydantic.total_candidates_considered,
-        processing_time_ms=api_response_pydantic.processing_time_ms,
-    )
+    final_mcp_responses = []
+
+    for response_pydantic in responses_list:
+        if not response_pydantic:
+            logger.warning("A backend search returned None; skipping this response.")
+            continue
+
+        actual_backend_results = response_pydantic.results
+        mcp_results_list = []
+        for backend_item in actual_backend_results[:tool_limit]:
+            mcp_results_list.append(_prepare_mcp_result_item(backend_item))
+
+        final_mcp_response = APISearchResponse(
+            query=response_pydantic.query,
+            packages_applied=response_pydantic.packages_applied,
+            results=mcp_results_list,
+            count=len(mcp_results_list),
+            total_candidates_considered=response_pydantic.total_candidates_considered,
+            processing_time_ms=response_pydantic.processing_time_ms,
+        )
+        final_mcp_responses.append(final_mcp_response.model_dump(exclude_none=True))
 
-    return final_mcp_response.model_dump(exclude_none=True)
+    return final_mcp_responses
 
 
 @mcp_app.tool()
-async def get_by_id(ctx: MCPContext, group_id: int) -> Optional[Dict[str, Any]]:
-    """Retrieves a specific statement group by its unique identifier.
+async def get_by_id(
+    ctx: MCPContext, group_id: Union[int, List[int]]
+) -> List[Optional[Dict[str, Any]]]:
+    """Retrieves specific statement groups by their unique identifier(s).
 
-    The `display_statement_text` field is omitted from the response.
+    The `display_statement_text` field is omitted from the response. This tool
+    always returns a list of results.
 
     Args:
         ctx: The MCP context, providing access to the backend service.
-        group_id: The unique integer identifier of the statement group to retrieve.
-                  For example, `12345`.
+        group_id: A single unique integer identifier or a list of identifiers
+                  of the statement group(s) to retrieve. For example, `12345` or
+                  `[12345, 67890]`.
 
     Returns:
-        A dictionary corresponding to the APISearchResultItem model if a
-        statement group with the given ID is found (with
-        `display_statement_text` omitted). Returns None (which will be
-        serialized as JSON null by MCP) if no such group exists.
+        A list of dictionaries, where each dictionary corresponds to the
+        APISearchResultItem model. If an ID is not found, its corresponding
+        entry in the list will be None (serialized as JSON null by MCP).
     """
     backend = await _get_backend_from_context(ctx)
-    logger.info(f"MCP Tool 'get_by_id' called for group_id: {group_id}")
+    logger.info(f"MCP Tool 'get_by_id' called for group_id(s): {group_id}")
 
-    backend_item: Optional[APISearchResultItem]
+    backend_items: Union[
+        Optional[APISearchResultItem], List[Optional[APISearchResultItem]]
+    ]
     if asyncio.iscoroutinefunction(backend.get_by_id):
-        backend_item = await backend.get_by_id(group_id=group_id)
+        backend_items = await backend.get_by_id(group_id=group_id)
     else:
-        backend_item = backend.get_by_id(group_id=group_id)
+        backend_items = backend.get_by_id(group_id=group_id)
+
+    # Normalize to a list for consistent return type
+    items_list = (
+        [backend_items] if not isinstance(backend_items, list) else backend_items
+    )
+
+    mcp_items = []
+    for item in items_list:
+        if item:
+            mcp_item = _prepare_mcp_result_item(item)
+            mcp_items.append(mcp_item.model_dump(exclude_none=True))
+        else:
+            mcp_items.append(None)
 
-    if backend_item:
-        mcp_item = _prepare_mcp_result_item(backend_item)
-        return mcp_item.model_dump(exclude_none=True)
-    return None
+    return mcp_items
 
 
 @mcp_app.tool()
-async def get_dependencies(ctx: MCPContext, group_id: int) -> Optional[Dict[str, Any]]:
-    """Retrieves the direct dependencies (citations) for a specific statement group.
+async def get_dependencies(
+    ctx: MCPContext, group_id: Union[int, List[int]]
+) -> List[Optional[Dict[str, Any]]]:
+    """Retrieves direct dependencies (citations) for specific statement group(s).
 
     The `display_statement_text` field within each cited item is omitted
-    from the response.
+    from the response. This tool always returns a list of results.
 
     Args:
         ctx: The MCP context, providing access to the backend service.
-        group_id: The unique integer identifier of the statement group for which
-                  to fetch its direct dependencies. For example, `12345`.
+        group_id: A single unique integer identifier or a list of identifiers for
+                  the statement group(s) for which to fetch direct dependencies.
+                  For example, `12345` or `[12345, 67890]`.
 
     Returns:
-        A dictionary corresponding to the APICitationsResponse model, which
-        contains a list of cited statement groups (each with
-        `display_statement_text` omitted), if the source group_id
-        is found and has dependencies. Returns None (serialized as JSON null
-        by MCP) if the source group is not found or has no dependencies.
+        A list of dictionaries, where each dictionary corresponds to the
+        APICitationsResponse model. If a source group ID is not found or has
+        no dependencies, its corresponding entry will be None.
     """
     backend = await _get_backend_from_context(ctx)
-    logger.info(f"MCP Tool 'get_dependencies' called for group_id: {group_id}")
+    logger.info(f"MCP Tool 'get_dependencies' called for group_id(s): {group_id}")
 
-    backend_response: Optional[APICitationsResponse]
+    backend_responses: Union[
+        Optional[APICitationsResponse], List[Optional[APICitationsResponse]]
+    ]
     if asyncio.iscoroutinefunction(backend.get_dependencies):
-        backend_response = await backend.get_dependencies(group_id=group_id)
+        backend_responses = await backend.get_dependencies(group_id=group_id)
     else:
-        backend_response = backend.get_dependencies(group_id=group_id)
+        backend_responses = backend.get_dependencies(group_id=group_id)
 
-    if backend_response:
-        mcp_citations_list = []
-        for backend_item in backend_response.citations:
-            mcp_citations_list.append(_prepare_mcp_result_item(backend_item))
-
-        final_mcp_response = APICitationsResponse(
-            source_group_id=backend_response.source_group_id,
-            citations=mcp_citations_list,
-            count=len(mcp_citations_list),
-        )
-        return final_mcp_response.model_dump(exclude_none=True)
-    return None
+    # Normalize to a list for consistent return type
+    responses_list = (
+        [backend_responses]
+        if not isinstance(backend_responses, list)
+        else backend_responses
+    )
+    final_mcp_responses = []
+
+    for response in responses_list:
+        if response:
+            mcp_citations_list = []
+            for backend_item in response.citations:
+                mcp_citations_list.append(_prepare_mcp_result_item(backend_item))
+
+            final_response = APICitationsResponse(
+                source_group_id=response.source_group_id,
+                citations=mcp_citations_list,
+                count=len(mcp_citations_list),
+            )
+            final_mcp_responses.append(final_response.model_dump(exclude_none=True))
+        else:
+            final_mcp_responses.append(None)
+
+    return final_mcp_responses

{lean_explore-0.2.1.dist-info → lean_explore-0.3.0.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: lean-explore
-Version: 0.2.1
-Summary: A project to explore and rank Lean mathematical declarations.
+Version: 0.3.0
+Summary: A search engine for Lean 4 declarations.
 Author-email: Justin Asher <justinchadwickasher@gmail.com>
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -213,15 +213,13 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Classifier: Topic :: Text Processing :: Indexing
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: sqlalchemy>=2.0
@@ -241,7 +239,28 @@ Requires-Dist: tqdm>=4.60
 Requires-Dist: requests>=2.25.0
 Dynamic: license-file
 
-# LeanExplore
+<h1 align="center">
+  LeanExplore
+</h1>
+
+<h3 align="center">
+  A search engine for Lean 4 declarations
+</h3>
+
+<p align="center">
+  <a href="https://pypi.org/project/lean-explore/">
+    <img src="https://img.shields.io/pypi/v/lean-explore.svg" alt="PyPI version" />
+  </a>
+  <a href="https://github.com/justincasher/lean-explore/blob/main/LeanExplore.pdf">
+    <img src="https://img.shields.io/badge/Paper-PDF-blue.svg" alt="Read the Paper" />
+  </a>
+  <a href="https://github.com/justincasher/lean-explore/commits/main">
+    <img src="https://img.shields.io/github/last-commit/justincasher/lean-explore" alt="last update" />
+  </a>
+  <a href="https://github.com/justincasher/lean-explore/blob/main/LICENSE">
+    <img src="https://img.shields.io/github/license/justincasher/lean-explore.svg" alt="license" />
+  </a>
+</p>
 
 A search engine for Lean 4 declarations. This project provides tools and resources for exploring the Lean 4 ecosystem.
 
@@ -251,6 +270,7 @@ The current indexed projects include:
 
 * Batteries
 * Lean
+* Init
 * Mathlib
 * PhysLean
 * Std

{lean_explore-0.2.1.dist-info → lean_explore-0.3.0.dist-info}/RECORD

@@ -1,26 +1,26 @@
 lean_explore/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
 lean_explore/defaults.py,sha256=IJw6od-y0grYbwiDJ5ewNZI4u0j0dCCu_AXCDwWLHuA,4459
 lean_explore/api/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
-lean_explore/api/client.py,sha256=AgZG7pUY53Tl1WOhJgUdT0yxa_O1sHsals0pnjRD-Pc,4839
+lean_explore/api/client.py,sha256=rvSIDbyqGl2I5b214VBBfT_UM9CvaHLa6DElsnUbi9E,7848
 lean_explore/cli/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
-lean_explore/cli/agent.py,sha256=jf1ebnViAqtKcZAGArqBf9YKHPbsOTpffOXr0Cd2M3Q,29933
+lean_explore/cli/agent.py,sha256=BaC5uoK5HrySBJCB0aGcgLOE1N-UYlmbWz2hUcNUk44,30509
 lean_explore/cli/config_utils.py,sha256=RyIaDNP1UpUQZoy7HfaZ_JOXUgtzUP51Zrq_s6q7urY,16639
 lean_explore/cli/data_commands.py,sha256=mTBqFU7-fF4ZBGzCmNawZA_eHy0jyEMLlBEDEBXpxwY,21462
 lean_explore/cli/main.py,sha256=ZdbXy8x2VQ--JARqJMa9iFnrOhOCLcVgjpWhXkxj80o,24323
 lean_explore/local/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
 lean_explore/local/search.py,sha256=ZW8rKJ2riT6RRi6ngo8SylxQ_5jQbsipuv84kqpiwc4,40930
-lean_explore/local/service.py,sha256=7VT_njCiKlMqQItJ-Uy2aDraLttulAAgPhRp2BMWnSk,16166
+lean_explore/local/service.py,sha256=AQAbYZ9tr3Yd_ED4weEnbRDwvkh7_0E-ERy1C1Abjlg,19292
 lean_explore/mcp/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
 lean_explore/mcp/app.py,sha256=XG6zTAaBRbdV1Ep_Za2JmifEmkYFKBmcGrdizCLlH-s,3808
 lean_explore/mcp/server.py,sha256=pzhLNGfTxelZvQ7ZJrWW0cbNH4MCwhviV24Y-yfQa0c,8666
-lean_explore/mcp/tools.py,sha256=Lri2GNIKNCLZNpNfvgvI600w3-0gaJGdPCFhhd7WVuk,9580
+lean_explore/mcp/tools.py,sha256=L1U76Xg1nh3mRzq_zuEltVE2R-rsm7m6i4DFPkhqS48,10263
 lean_explore/shared/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
 lean_explore/shared/models/__init__.py,sha256=LK4g9wj7jCilTUfQcdQAxDf3F2GjMwzQZJYgnQ8ciGo,38
 lean_explore/shared/models/api.py,sha256=jejNDpgj-cu0KZTqkuOjM0useN4EvhvNB19lFFAOV94,4635
 lean_explore/shared/models/db.py,sha256=JYfIBnPrHZO2j7gHAVMlw9WSqVC2NinCG5KuBzdQWyk,16099
-lean_explore-0.2.1.dist-info/licenses/LICENSE,sha256=l4QLw1kIvEOjUktmmKm4dycK1E249Qs2s2AQTYbMXpY,11354
-lean_explore-0.2.1.dist-info/METADATA,sha256=5EWx5NniczmS6ApKVvoHj1RfgNC6eO9JgOIyZyNA1SY,15611
-lean_explore-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lean_explore-0.2.1.dist-info/entry_points.txt,sha256=JXl2Mo3BRX4jAU-Nxg_CWJR790pB_oi5qnt3Pv5iZnk,58
-lean_explore-0.2.1.dist-info/top_level.txt,sha256=h51BKWrFvB7iym-IlaNAAHX5MZfA8Gmg-aDuXGo0fQ8,13
-lean_explore-0.2.1.dist-info/RECORD,,
+lean_explore-0.3.0.dist-info/licenses/LICENSE,sha256=l4QLw1kIvEOjUktmmKm4dycK1E249Qs2s2AQTYbMXpY,11354
+lean_explore-0.3.0.dist-info/METADATA,sha256=gJTIosn6cuK8s1x0Q8XD4s5RZydzow_jUZFlZsKUpIM,16304
+lean_explore-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lean_explore-0.3.0.dist-info/entry_points.txt,sha256=JXl2Mo3BRX4jAU-Nxg_CWJR790pB_oi5qnt3Pv5iZnk,58
+lean_explore-0.3.0.dist-info/top_level.txt,sha256=h51BKWrFvB7iym-IlaNAAHX5MZfA8Gmg-aDuXGo0fQ8,13
+lean_explore-0.3.0.dist-info/RECORD,,