From 84d56fbb313c65d67995acb0775f1eba8ec881d2 Mon Sep 17 00:00:00 2001 From: Benoit Kohler Date: Mon, 18 Aug 2025 11:16:16 +0200 Subject: [PATCH] Stable to Develop (#14) Stable to Develop (Cleanup) --- .../python-docstring.instructions.md | 15 ++ .github/instructions/python.instructions.md | 25 ++ .github/instructions/tooling.instructions.md | 11 + .github/workflows/ci.yml | 3 +- .github/workflows/sync-docs.yml | 68 +++--- LICENSE.txt | 201 +++++++++++++++ README.md | 4 + pyproject.toml | 7 +- src/infrahub_mcp/__init__.py | 0 src/infrahub_mcp/branch.py | 63 +++++ src/infrahub_mcp/constants.py | 9 + src/infrahub_mcp/gql.py | 44 ++++ src/infrahub_mcp/nodes.py | 230 ++++++++++++++++++ src/infrahub_mcp/prompts/main.md | 8 + src/infrahub_mcp/schema.py | 140 +++++++++++ src/infrahub_mcp/server.py | 35 +++ src/infrahub_mcp/utils.py | 94 +++++++ tests/integration/conftest.py | 5 +- tests/unit/test_tools.py | 2 +- uv.lock | 2 +- 20 files changed, 924 insertions(+), 42 deletions(-) create mode 100644 .github/instructions/python-docstring.instructions.md create mode 100644 .github/instructions/python.instructions.md create mode 100644 .github/instructions/tooling.instructions.md create mode 100644 LICENSE.txt create mode 100644 src/infrahub_mcp/__init__.py create mode 100644 src/infrahub_mcp/branch.py create mode 100644 src/infrahub_mcp/constants.py create mode 100644 src/infrahub_mcp/gql.py create mode 100644 src/infrahub_mcp/nodes.py create mode 100644 src/infrahub_mcp/prompts/main.md create mode 100644 src/infrahub_mcp/schema.py create mode 100644 src/infrahub_mcp/server.py create mode 100644 src/infrahub_mcp/utils.py diff --git a/.github/instructions/python-docstring.instructions.md b/.github/instructions/python-docstring.instructions.md new file mode 100644 index 0000000..fce61b7 --- /dev/null +++ b/.github/instructions/python-docstring.instructions.md @@ -0,0 +1,15 @@ +--- +applyTo: '**/*.py' +--- + +Follow these guidelines for writing clear and consistent Python docstrings: + +- Always use triple quotes (""") for docstrings +- Follow Google-style docstring format +- Include these sections when applicable: + - Brief one-line description + - Detailed description (if needed) + - Args/Parameters without typing + - Returns + - Raises + - Examples diff --git a/.github/instructions/python.instructions.md b/.github/instructions/python.instructions.md new file mode 100644 index 0000000..af729b0 --- /dev/null +++ b/.github/instructions/python.instructions.md @@ -0,0 +1,25 @@ +--- +applyTo: '**/*.py' +--- + +# Python Coding Standards + +- Use Python 3.10+ syntax and features. +- Follow PEP 8 for code style and formatting. +- Use type annotations for all function signatures and variables. +- Prefer f-strings for string formatting. +- Avoid wildcard imports; import only what you need. +- Organize imports using `ruff` or `isort` conventions: standard library, third-party, local. +- Write small, single-responsibility functions and classes. +- Avoid global variables; prefer dependency injection or class attributes. +- Use context managers (`with` statements) for resource management. +- Prefer list comprehensions and generator expressions over manual loops when appropriate. +- Use built-in exceptions or define custom exceptions for error handling. +- Document all public modules, classes, and functions. + +## Formatting and Linting + +The project is using Use ruff and mypy for type checking and validations. + +you can format all python files by running `uv run invoke format` +and you can validate that all files are formatted correctly by running `uv run invoke lint` diff --git a/.github/instructions/tooling.instructions.md b/.github/instructions/tooling.instructions.md new file mode 100644 index 0000000..0816923 --- /dev/null +++ b/.github/instructions/tooling.instructions.md @@ -0,0 +1,11 @@ +--- +applyTo: '*' +--- + +# Tooling rules + +- Prefer Github Actions for automated validation when user commit to the repository +- Don't use Git precommit +- Use ruff and mypy to validate and lint python files +- Use yamllint to validate yaml files +- Use uv to manage the python project and its dependencies \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 68ae02a..86aa08d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,7 +5,8 @@ on: pull_request: push: branches: - - main + - stable + - develop concurrency: group: ${{ github.workflow }}-${{ github.ref }} diff --git a/.github/workflows/sync-docs.yml b/.github/workflows/sync-docs.yml index 4e34694..288c9e8 100644 --- a/.github/workflows/sync-docs.yml +++ b/.github/workflows/sync-docs.yml @@ -1,39 +1,39 @@ ---- +# --- # yamllint disable rule:truthy -name: Sync Folders +# name: Sync Folders -on: - push: - branches: - - main - paths: - - 'docs/docs/**' - - 'docs/sidebars.ts' +# on: +# push: +# branches: +# - stable +# paths: +# - 'docs/docs/**' +# - 'docs/sidebars.ts' -jobs: - sync: - runs-on: ubuntu-latest - steps: - - name: Checkout source repository - uses: actions/checkout@v4 - with: - path: source-repo +# jobs: +# sync: +# runs-on: ubuntu-latest +# steps: +# - name: Checkout source repository +# uses: actions/checkout@v4 +# with: +# path: source-repo - - name: Checkout target repository - uses: actions/checkout@v4 - with: - repository: opsmill/infrahub-docs - token: ${{ secrets.PAT_TOKEN }} - path: target-repo +# - name: Checkout target repository +# uses: actions/checkout@v4 +# with: +# repository: opsmill/infrahub-docs +# token: ${{ secrets.PAT_TOKEN }} +# path: target-repo - - name: Sync folders - run: | - rm -rf target-repo/docs/docs-mcp/* - rm -f target-repo/docs/sidebars-mcp.ts - cp -r source-repo/docs/docs/* target-repo/docs/docs-mcp/ - cp source-repo/docs/sidebars.ts target-repo/docs/sidebars-mcp.ts - cd target-repo - git config user.name github-actions - git config user.email github-actions@github.com - git add . - if ! (git diff --quiet && git diff --staged --quiet); then git commit -m "Sync docs from mcp repo" && git push; fi +# - name: Sync folders +# run: | +# rm -rf target-repo/docs/docs-mcp/* +# rm -f target-repo/docs/sidebars-mcp.ts +# cp -r source-repo/docs/docs/* target-repo/docs/docs-mcp/ +# cp source-repo/docs/sidebars.ts target-repo/docs/sidebars-mcp.ts +# cd target-repo +# git config user.name github-actions +# git config user.email github-actions@github.com +# git add . +# if ! (git diff --quiet && git diff --staged --quiet); then git commit -m "Sync docs from mcp repo" && git push; fi diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000..06154b9 --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2024 OpsMill SAS + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/README.md b/README.md index 8e2884c..fae42d9 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ MCP server to interact with Infrahub - fastmcp - infrahub_sdk + ## Installation 1. **Clone the repo** @@ -29,6 +30,7 @@ MCP server to interact with Infrahub uv run fastmcp run src/infrahub_mcp_server/server.py:mcp ``` + ## Configuration Set the following environment variables as needed: @@ -40,6 +42,7 @@ Set the following environment variables as needed: | `MCP_HOST` | Host for the web server | `0.0.0.0` | | `MCP_PORT` | Port for the web server | `8001` | + ## Usage ### HTTP API mode @@ -144,3 +147,4 @@ Response (all schemas): "schemas": [ { ... }, ... ] } ``` + diff --git a/pyproject.toml b/pyproject.toml index 4b7e237..aad6cc0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] -name = "infrahub-mcp-server" +name = "infrahub-mcp" version = "0.1.0" -description = "Add your description here" +description = "MCP server to interact with Infrahub" readme = "README.md" requires-python = ">=3.13" dependencies = [ @@ -85,7 +85,8 @@ skip-magic-trailing-comma = false line-ending = "auto" [tool.ruff.lint.isort] -known-first-party = ["infrahub_mcp_server"] +known-first-party = ["infrahub_mcp"] + [tool.ruff.lint.pycodestyle] max-line-length = 150 diff --git a/src/infrahub_mcp/__init__.py b/src/infrahub_mcp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/infrahub_mcp/branch.py b/src/infrahub_mcp/branch.py new file mode 100644 index 0000000..e289d8e --- /dev/null +++ b/src/infrahub_mcp/branch.py @@ -0,0 +1,63 @@ +from typing import TYPE_CHECKING, Annotated + +from fastmcp import Context, FastMCP +from infrahub_sdk.branch import BranchData +from infrahub_sdk.exceptions import GraphQLError +from mcp.types import ToolAnnotations +from pydantic import Field + +from infrahub_mcp.utils import MCPResponse, MCPToolStatus, _log_and_return_error + +if TYPE_CHECKING: + from infrahub_sdk import InfrahubClient + +mcp: FastMCP = FastMCP(name="Infrahub Branches") + + +@mcp.tool( + tags=["branches", "create"], + annotations=ToolAnnotations(readOnlyHint=False, idempotentHint=True, destructiveHint=False), +) +async def branch_create( + ctx: Context, + name: Annotated[str, Field(description="Name of the branch to create.")], + sync_with_git: Annotated[bool, Field(default=False, description="Whether to sync the branch with git.")], +) -> MCPResponse[dict[str, str]]: + """Create a new branch in infrahub. + + Parameters: + name: Name of the branch to create. + sync_with_git: Whether to sync the branch with git. Defaults to False. + + Returns: + Dictionary with success status and branch details. + """ + + client: InfrahubClient = ctx.request_context.lifespan_context.client + ctx.info(f"Creating branch {name} in Infrahub...") + + try: + branch = await client.branch.create(branch_name=name, sync_with_git=sync_with_git, background_execution=False) + + except GraphQLError as exc: + return _log_and_return_error(ctx=ctx, error=exc, remediation="Check the branch name or your permissions.") + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data={ + "name": branch.name, + "id": branch.id, + }, + ) + + +@mcp.tool(tags=["branches", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_branches(ctx: Context) -> MCPResponse[dict[str, BranchData]]: + """Retrieve all branches from infrahub.""" + + client: InfrahubClient = ctx.request_context.lifespan_context.client + ctx.info("Fetching all branches from Infrahub...") + + branches: dict[str, BranchData] = await client.branch.all() + + return MCPResponse(status=MCPToolStatus.SUCCESS, data=branches) diff --git a/src/infrahub_mcp/constants.py b/src/infrahub_mcp/constants.py new file mode 100644 index 0000000..01dde86 --- /dev/null +++ b/src/infrahub_mcp/constants.py @@ -0,0 +1,9 @@ +NAMESPACES_INTERNAL = ["Internal", "Profile", "Template"] + +schema_attribute_type_mapping = { + "Text": "String", + "Number": "Integer", + "Boolean": "Boolean", + "DateTime": "DateTime", + "Enum": "String", +} diff --git a/src/infrahub_mcp/gql.py b/src/infrahub_mcp/gql.py new file mode 100644 index 0000000..f890032 --- /dev/null +++ b/src/infrahub_mcp/gql.py @@ -0,0 +1,44 @@ +from typing import TYPE_CHECKING, Annotated, Any + +from fastmcp import Context, FastMCP +from mcp.types import ToolAnnotations +from pydantic import Field + +from infrahub_mcp.utils import MCPResponse, MCPToolStatus + +if TYPE_CHECKING: + from infrahub_sdk import InfrahubClient + +mcp: FastMCP = FastMCP(name="Infrahub GraphQL") + + +@mcp.tool(tags=["schemas", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_graphql_schema(ctx: Context) -> MCPResponse[str]: + """Retrieve the GraphQL schema from Infrahub + + Parameters: + None + + Returns: + MCPResponse with the GraphQL schema as a string. + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + resp = await client._get(url=f"{client.address}/schema.graphql") # noqa: SLF001 + return MCPResponse(status=MCPToolStatus.SUCCESS, data=resp.text) + + +@mcp.tool(tags=["schemas", "retrieve"], annotations=ToolAnnotations(readOnlyHint=False)) +async def query_graphql( + ctx: Context, query: Annotated[str, Field(description="GraphQL query to execute.")] +) -> MCPResponse[dict[str, Any]]: + """Execute a GraphQL query against Infrahub. + + Parameters: + query: GraphQL query to execute. + + Returns: + MCPResponse with the result of the query. + + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + return await client.execute_graphql(query=query) diff --git a/src/infrahub_mcp/nodes.py b/src/infrahub_mcp/nodes.py new file mode 100644 index 0000000..b6d0c07 --- /dev/null +++ b/src/infrahub_mcp/nodes.py @@ -0,0 +1,230 @@ +from typing import TYPE_CHECKING, Annotated, Any + +from fastmcp import Context, FastMCP +from infrahub_sdk.exceptions import GraphQLError, SchemaNotFoundError +from infrahub_sdk.types import Order +from mcp.types import ToolAnnotations +from pydantic import Field + +from infrahub_mcp.constants import schema_attribute_type_mapping +from infrahub_mcp.utils import MCPResponse, MCPToolStatus, _log_and_return_error, convert_node_to_dict + +if TYPE_CHECKING: + from infrahub_sdk.client import InfrahubClient + +mcp: FastMCP = FastMCP(name="Infrahub Nodes") + + +@mcp.tool(tags=["nodes", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_nodes( + ctx: Context, + kind: Annotated[str, Field(description="Kind of the objects to retrieve.")], + branch: Annotated[ + str | None, + Field(default=None, description="Branch to retrieve the objects from. Defaults to None (uses default branch)."), + ], + filters: Annotated[dict[str, Any] | None, Field(default=None, description="Dictionary of filters to apply.")], + partial_match: Annotated[bool, Field(default=False, description="Whether to use partial matching for filters.")], +) -> MCPResponse[list[str]]: + """Get all objects of a specific kind from Infrahub. + + To retrieve the list of available kinds, use the `get_schema_mapping` tool. + To retrieve the list of available filters for a specific kind, use the `get_node_filters` tool. + + Parameters: + kind: Kind of the objects to retrieve. + branch: Branch to retrieve the objects from. Defaults to None (uses default branch). + filters: Dictionary of filters to apply. + partial_match: Whether to use partial matching for filters. + + Returns: + MCPResponse with success status and objects. + + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + ctx.info(f"Fetching nodes of kind: {kind} with filters: {filters} from Infrahub...") + + # Verify if the kind exists in the schema and guide Tool if not + try: + schema = await client.schema.get(kind=kind, branch=branch) + except SchemaNotFoundError: + error_msg = f"Schema not found for kind: {kind}." + remediation_msg = "Use the `get_schema_mapping` tool to list available kinds." + return _log_and_return_error(ctx=ctx, error=error_msg, remediation=remediation_msg) + + # TODO: Verify if the filters are valid for the kind and guide Tool if not + + try: + if filters: + nodes = await client.filters( + kind=schema.kind, + branch=branch, + partial_match=partial_match, + parallel=True, + order=Order(disable=True), + populate_store=True, + prefetch_relationships=True, + **filters, + ) + else: + nodes = await client.all( + kind=schema.kind, + branch=branch, + parallel=True, + order=Order(disable=True), + populate_store=True, + prefetch_relationships=True, + ) + except GraphQLError as exc: + return _log_and_return_error(ctx=ctx, error=exc, remediation="Check the provided filters or the kind name.") + + # Format the response with serializable data + # serialized_nodes = [] + # for node in nodes: + # node_data = await convert_node_to_dict(obj=node, branch=branch) + # serialized_nodes.append(node_data) + serialized_nodes = [obj.display_label for obj in nodes] + + # Return the serialized response + ctx.debug(f"Retrieved {len(serialized_nodes)} nodes of kind {kind}") + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data=serialized_nodes, + ) + + +@mcp.tool(tags=["nodes", "filters", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_node_filters( + ctx: Context, + kind: Annotated[str, Field(description="Kind of the objects to retrieve.")], + branch: Annotated[ + str | None, + Field(default=None, description="Branch to retrieve the objects from. Defaults to None (uses default branch)."), + ], +) -> MCPResponse[dict[str, str]]: + """Retrieve all the available filters for a specific schema node kind. + + There's multiple types of filters + attribute filters are in the form attribute__value + + relationship filters are in the form relationship__attribute__value + you can find more information on the peer node of the relationship using the `get_schema` tool + + Filters that start with parent refer to a related generic schema node. + You can find the type of that related node by inspected the output of the `get_schema` tool. + + Parameters: + kind: Kind of the objects to retrieve. + branch: Branch to retrieve the objects from. Defaults to None (uses default branch). + + Returns: + MCPResponse with success status and filters. + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + ctx.info(f"Fetching available filters for kind: {kind} from Infrahub...") + + # Verify if the kind exists in the schema and guide Tool if not + try: + schema = await client.schema.get(kind=kind, branch=branch) + except SchemaNotFoundError: + error_msg = f"Schema not found for kind: {kind}." + remediation_msg = "Use the `get_schema_mapping` tool to list available kinds." + return _log_and_return_error(ctx=ctx, error=error_msg, remediation=remediation_msg) + + filters = { + f"{attribute.name}__value": schema_attribute_type_mapping.get(attribute.kind, "String") + for attribute in schema.attributes + } + + for relationship in schema.relationships: + relationship_schema = await client.schema.get(kind=relationship.peer) + relationship_filters = { + f"{relationship.name}__{attribute.name}__value": schema_attribute_type_mapping.get(attribute.kind, "String") + for attribute in relationship_schema.attributes + } + filters.update(relationship_filters) + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data=filters, + ) + + +@mcp.tool(tags=["nodes", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_related_nodes( + ctx: Context, + kind: Annotated[str, Field(description="Kind of the objects to retrieve.")], + relation: Annotated[str, Field(description="Name of the relation to fetch.")], + filters: Annotated[dict[str, Any] | None, Field(default=None, description="Dictionary of filters to apply.")], + branch: Annotated[ + str | None, + Field(default=None, description="Branch to retrieve the objects from. Defaults to None (uses default branch)."), + ], +) -> MCPResponse[list[dict[str, Any]]]: + """Retrieve related nodes by relation name and a kind. + + Args: + kind: Kind of the node to fetch. + filters: Filters to apply on the node to fetch. + relation: Name of the relation to fetch. + branch: Branch to fetch the node from. Defaults to None (uses default branch). + + Returns: + MCPResponse with success status and objects. + + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + filters = filters or {} + if branch: + ctx.info(f"Fetching nodes related to {kind} with filters {filters} in branch {branch} from Infrahub...") + else: + ctx.info(f"Fetching nodes related to {kind} with filters {filters} from Infrahub...") + + try: + node_id = node_hfid = None + if filters.get("ids"): + node_id = filters["ids"][0] + elif filters.get("hfid"): + node_hfid = filters["hfid"] + if node_id: + node = await client.get( + kind=kind, + id=node_id, + branch=branch, + include=[relation], + prefetch_relationships=True, + populate_store=True, + ) + elif node_hfid: + node = await client.get( + kind=kind, + hfid=node_hfid, + branch=branch, + include=[relation], + prefetch_relationships=True, + populate_store=True, + ) + except Exception as exc: # noqa: BLE001 + return _log_and_return_error(exc) + + rel = getattr(node, relation, None) + if not rel: + _log_and_return_error( + ctx=ctx, + error=f"Relation '{relation}' not found in kind '{kind}'.", + remediation="Check the schema for the kind to confirm if the relation exists.", + ) + peers = [ + await convert_node_to_dict( + branch=branch, + obj=peer.peer, + include_id=True, + ) + for peer in rel.peers + ] + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data=peers, + ) diff --git a/src/infrahub_mcp/prompts/main.md b/src/infrahub_mcp/prompts/main.md new file mode 100644 index 0000000..7218907 --- /dev/null +++ b/src/infrahub_mcp/prompts/main.md @@ -0,0 +1,8 @@ +You are an infrastructure specilist specialized in answering questions about the infrastructure. + +All the information you need are present in Infrahub and you can access it via an MCP server which exposes a number of tools. + +When someone ask about a specific data, you need to: +- Identify what is the associated kind in the schema for this data using the tool `schema_get_mapping` +- Retrieve more information about this specific model, including the option available to filter (tool : `get_node_filters`) +- Use the tool `get_objects` to query one or multiple objects \ No newline at end of file diff --git a/src/infrahub_mcp/schema.py b/src/infrahub_mcp/schema.py new file mode 100644 index 0000000..a911fe0 --- /dev/null +++ b/src/infrahub_mcp/schema.py @@ -0,0 +1,140 @@ +from typing import TYPE_CHECKING, Annotated, Any + +from fastmcp import Context, FastMCP +from infrahub_sdk.exceptions import BranchNotFoundError, SchemaNotFoundError +from mcp.types import ToolAnnotations +from pydantic import Field + +from infrahub_mcp.constants import NAMESPACES_INTERNAL +from infrahub_mcp.utils import MCPResponse, MCPToolStatus, _log_and_return_error + +if TYPE_CHECKING: + from infrahub_sdk import InfrahubClient + +mcp: FastMCP = FastMCP(name="Infrahub Schemas") + + +@mcp.tool(tags=["schemas", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_schema_mapping( + ctx: Context, + branch: Annotated[ + str | None, + Field(default=None, description="Branch to retrieve the mapping from. Defaults to None (uses default branch)."), + ], +) -> MCPResponse[dict[str, str]]: + """List all schema nodes and generics available in Infrahub + + Parameters: + branch: Branch to retrieve the mapping from. Defaults to None (uses default branch). + + Returns: + Dictionary with success status and schema mapping. + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + if branch: + ctx.info(f"Fetching schema mapping for {branch} from Infrahub...") + else: + ctx.info("Fetching schema mapping from Infrahub...") + + try: + all_schemas = await client.schema.all(branch=branch) + except BranchNotFoundError as exc: + return _log_and_return_error(ctx=ctx, error=exc, remediation="Check the branch name or your permissions.") + + # TODO: Should we add the description ? + schema_mapping = { + kind: node.label or "" for kind, node in all_schemas.items() if node.namespace not in NAMESPACES_INTERNAL + } + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data=schema_mapping, + ) + + +@mcp.tool(tags=["schemas", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_schema( + ctx: Context, + kind: Annotated[str, Field(description="Schema Kind to retrieve.")], + branch: Annotated[ + str | None, + Field(default=None, description="Branch to retrieve the schema from. Defaults to None (uses default branch)."), + ], +) -> MCPResponse[dict[str, Any]]: + """Retrieve the full schema for a specific kind. + This includes attributes, relationships, and their types. + + Parameters: + kind: Schema Kind to retrieve. + branch: Branch to retrieve the schema from. Defaults to None (uses default branch). + + Returns: + Dictionary with success status and schema. + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + ctx.info(f"Fetching schema of {kind} from Infrahub...") + + try: + schema = await client.schema.get(kind=kind, branch=branch) + except SchemaNotFoundError: + error_msg = f"Schema not found for kind: {kind}." + remediation_msg = "Use the `get_schema_mapping` tool to list available kinds." + return _log_and_return_error(ctx=ctx, error=error_msg, remediation=remediation_msg) + except BranchNotFoundError as exc: + return _log_and_return_error(ctx=ctx, error=exc, remediation="Check the branch name or your permissions.") + + schema = await client.schema.get(kind=kind, branch=branch) + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data=schema.model_dump(), + ) + + +@mcp.tool(tags=["schemas", "retrieve"], annotations=ToolAnnotations(readOnlyHint=True)) +async def get_schemas( + ctx: Context, + branch: Annotated[ + str | None, + Field(default=None, description="Branch to retrieve schemas from. Defaults to None (uses default branch)."), + ], + exclude_profiles: Annotated[ + bool, Field(default=True, description="Whether to exclude Profile schemas. Defaults to True.") + ], + exclude_templates: Annotated[ + bool, Field(default=True, description="Whether to exclude Template schemas. Defaults to True.") + ], +) -> MCPResponse[dict[str, dict[str, Any]]]: + """Retrieve all schemas from Infrahub, optionally excluding Profiles and Templates. + + Parameters: + infrahub_client: Infrahub client to use + branch: Branch to retrieve schemas from + exclude_profiles: Whether to exclude Profile schemas. Defaults to True. + exclude_templates: Whether to exclude Template schemas. Defaults to True. + + Returns: + Dictionary with success status and schemas. + + """ + client: InfrahubClient = ctx.request_context.lifespan_context.client + ctx.info(f"Fetching all schemas in branch {branch or 'main'} from Infrahub...") + + try: + all_schemas = await client.schema.all(branch=branch) + except BranchNotFoundError as exc: + return _log_and_return_error(ctx=ctx, error=exc, remediation="Check the branch name or your permissions.") + + # Filter out Profile and Template if requested + filtered_schemas = {} + for kind, schema in all_schemas.items(): + if (exclude_templates and schema.namespace == "Template") or ( + exclude_profiles and schema.namespace == "Profile" + ): + continue + filtered_schemas[kind] = schema.model_dump() + + return MCPResponse( + status=MCPToolStatus.SUCCESS, + data=filtered_schemas, + ) diff --git a/src/infrahub_mcp/server.py b/src/infrahub_mcp/server.py new file mode 100644 index 0000000..35486a7 --- /dev/null +++ b/src/infrahub_mcp/server.py @@ -0,0 +1,35 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass + +from fastmcp import FastMCP +from infrahub_sdk.client import InfrahubClient + +from infrahub_mcp.branch import mcp as branch_mcp +from infrahub_mcp.gql import mcp as graphql_mcp +from infrahub_mcp.nodes import mcp as nodes_mcp +from infrahub_mcp.schema import mcp as schema_mcp + + +@dataclass +class AppContext: + client: InfrahubClient + + +@asynccontextmanager +async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]: # noqa: ARG001, RUF029 + """Manages application lifecycle with type-safe context for the FastMCP server.""" + client = InfrahubClient() + try: + yield AppContext(client=client) + finally: + pass + + +mcp: FastMCP = FastMCP(name="Infrahub MCP Server", version="0.1.0", lifespan=app_lifespan) + +# Mount the various MCPs to the main server +mcp.mount(branch_mcp) +mcp.mount(graphql_mcp) +mcp.mount(nodes_mcp) +mcp.mount(schema_mcp) diff --git a/src/infrahub_mcp/utils.py b/src/infrahub_mcp/utils.py new file mode 100644 index 0000000..35b27e6 --- /dev/null +++ b/src/infrahub_mcp/utils.py @@ -0,0 +1,94 @@ +from enum import Enum +from pathlib import Path +from typing import Any, TypeVar + +from fastmcp import Context +from infrahub_sdk.node import Attribute, InfrahubNode, RelatedNode, RelationshipManager +from pydantic import BaseModel + +CURRENT_DIRECTORY = Path(__file__).parent.resolve() +PROMPTS_DIRECTORY = CURRENT_DIRECTORY / "prompts" + + +T = TypeVar("T") + + +class MCPToolStatus(Enum): + SUCCESS = "success" + ERROR = "error" + + +class MCPResponse[T](BaseModel): + status: MCPToolStatus + data: T | None = None + error: str | None = None + remediation: str | None = None + + +def get_prompt(name: str) -> str: + prompt_file = PROMPTS_DIRECTORY / f"{name}.md" + if not prompt_file.exists(): + raise FileNotFoundError(f"Prompt file '{prompt_file}' does not exist.") + return (PROMPTS_DIRECTORY / f"{name}.md").read_text() + + +def _log_and_return_error(ctx: Context, error: str | Exception, remediation: str | None = None) -> MCPResponse: + """Log an error and return a standardized error response.""" + if isinstance(error, Exception): + error = str(error) + ctx.error(message=error) + return MCPResponse( + status=MCPToolStatus.ERROR, + error=error, + remediation=remediation, + ) + + +async def convert_node_to_dict(*, obj: InfrahubNode, branch: str | None, include_id: bool = True) -> dict[str, Any]: # noqa: C901 + data = {} + + if include_id: + data["index"] = obj.id or None + + for attr_name in obj._schema.attribute_names: # noqa: SLF001 + attr: Attribute = getattr(obj, attr_name) + data[attr_name] = str(attr.value) + + for rel_name in obj._schema.relationship_names: # noqa: SLF001 + rel = getattr(obj, rel_name) + if rel and isinstance(rel, RelatedNode): + if not rel.initialized: + await rel.fetch() + related_node = obj._client.store.get( # noqa: SLF001 + branch=branch, + key=rel.peer.id, + raise_when_missing=False, + ) + if related_node: + data[rel_name] = ( + related_node.get_human_friendly_id_as_string(include_kind=True) + if related_node.hfid + else related_node.id + ) + elif rel and isinstance(rel, RelationshipManager): + peers: list[dict[str, Any]] = [] + if not rel.initialized: + await rel.fetch() + for peer in rel.peers: + # FIXME: We are using the store to avoid doing to many queries to Infrahub + # but we could end up doing store+infrahub if the store is not populated + related_node = obj._client.store.get( # noqa: SLF001 + key=peer.id, + raise_when_missing=False, + branch=branch, + ) + if not related_node: + await peer.fetch() + related_node = peer.peer + peers.append( + related_node.get_human_friendly_id_as_string(include_kind=True) + if related_node.hfid + else related_node.id, + ) + data[rel_name] = peers + return data diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 50f4218..4edb133 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -3,7 +3,8 @@ import pytest from agents.mcp import MCPServerStdio, MCPServerStdioParams -from infrahub_mcp_server.utils import get_prompt +from infrahub_mcp.utils import get_prompt + CURRENT_DIRECTORY = Path(__file__).parent.resolve() ROOT_DIRECTORY = CURRENT_DIRECTORY.parent.parent.resolve() @@ -30,7 +31,7 @@ def local_mcp_server() -> MCPServerStdio: "fastmcp", "run", "--no-banner", - "src/infrahub_mcp_server/server.py:mcp", + "src/infrahub_mcp/server.py:mcp", ], env={ "INFRAHUB_ADDRESS": "https://sandbox.infrahub.app", diff --git a/tests/unit/test_tools.py b/tests/unit/test_tools.py index 0679a7b..787e600 100644 --- a/tests/unit/test_tools.py +++ b/tests/unit/test_tools.py @@ -1,6 +1,6 @@ from fastmcp import Client -from infrahub_mcp_server.server import mcp +from infrahub_mcp.server import mcp async def test_list_schema() -> None: diff --git a/uv.lock b/uv.lock index 9e50071..c10f834 100644 --- a/uv.lock +++ b/uv.lock @@ -591,7 +591,7 @@ wheels = [ ] [[package]] -name = "infrahub-mcp-server" +name = "infrahub-mcp" version = "0.1.0" source = { editable = "." } dependencies = [