Astra MCP Server

#!/usr/bin/env python3 """ Tool Specification Generator for Astra DB Tables This program connects to Astra DB, retrieves table schema and sample data, and generates a tool specification in the format compatible with tool_config.json. """ import argparse import json import os import sys from typing import Any, Dict, List, Optional from dotenv import load_dotenv from pydantic_core.core_schema import none_schema from astrapy.info import TableIndexType # Add the parent directory to the path to import agentic_astra modules sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) from agentic_astra.database import AstraDBManager from agentic_astra.logger import get_logger from agentic_astra.llm import run_prompt from agentic_astra.tool_agent_prompt import prompt as tool_agent_prompt # Load environment variables load_dotenv(override=True) logger = get_logger("Agent Tool Generator", stdout=True) def export_prompt(file_path: str) -> str: prompt = tool_agent_prompt with open(file_path, "w") as f: f.write(prompt) logger.info(f"Prompt exported to {file_path}") class AstraToolAgent: db_manager = None db = None table = None collection = None table_name = None collection_name = None keyspace_name = None metadata = {} def __init__(self, astra_db_token: str, astra_db_name: str, table_name: str = None, collection_name: str = None, keyspace_name: str = "default_keyspace") -> Optional[Dict[str, Any]]: """Initialize the tool agent.""" self.db_manager = AstraDBManager(token=astra_db_token, db_name=astra_db_name) self.db = self.db_manager.get_db_by_name(astra_db_name) self.keyspace_name = keyspace_name if table_name: self.table_name = table_name self.table = self.db.get_table(self.table_name) def get_indexed_columns(self, table_indexes: List[Dict[str, Any]], index_type: TableIndexType = TableIndexType.REGULAR) -> List[str]: indexed_columns = [] for index in table_indexes: if index.index_type == index_type: indexed_columns.append(index.definition.column) return indexed_columns def get_table_schema(self) -> Optional[Dict[str, Any]]: """Get table schema information from Astra DB.""" tables = self.db.list_tables() table_metadata = next((table for table in tables if table.name == self.table_name), None) table_indexes = self.table.list_indexes() metadata = { "table_name": self.table_name, "keyspace_name": self.keyspace_name, "columns": table_metadata.raw_descriptor["definition"]["columns"], "primary_key": table_metadata.raw_descriptor["definition"]["primaryKey"], "partition_key": table_metadata.raw_descriptor["definition"]["primaryKey"]["partitionBy"], "partition_sort": table_metadata.raw_descriptor["definition"]["primaryKey"]["partitionSort"], "indexed_columns": self.get_indexed_columns(table_indexes, TableIndexType.REGULAR), "vector_columns": self.get_indexed_columns(table_indexes, TableIndexType.VECTOR), "text_columns": self.get_indexed_columns(table_indexes, TableIndexType.TEXT) } logger.info(f"Table metadata: {metadata}") logger.info(f"Retrieved schema for table '{self.table_name}': {metadata}") return metadata def get_sample_records(self, limit: int = 5) -> Optional[List[Dict[str, Any]]]: """Get sample records from the table.""" try: # Get first few records result = self.table.find(limit=limit) records = list(result) logger.info(f"Retrieved {len(records)} sample records from table '{self.table.name}'") return records except Exception as e: logger.error(f"Failed to get sample records from table '{self.table.name}': {str(e)}") return None def load_prompt_template(self, prompt_file: str = None) -> str: """Load the prompt template from the markdown file.""" if prompt_file: prompt_file_path = prompt_file else: prompt_file_path = os.path.join(os.path.dirname(__file__), "tool_specification_prompt.md") try: with open(prompt_file_path, 'r', encoding='utf-8') as f: return f.read() except FileNotFoundError: logger.error(f"Prompt template file not found: {prompt_file_path}") raise except Exception as e: logger.error(f"Error reading prompt template file: {e}") raise def generate_tool_specification(self, sample_size: int = 5, additional_instructions: str = "", prompt_file: str = None) -> Dict[str, Any]: """Generate tool specification based on table schema and sample data.""" self.metadata = self.get_table_schema() sample_records = self.get_sample_records(limit=sample_size) # Load prompt template from markdown file and format with variables prompt_template = "" if prompt_file: prompt_template = self.load_prompt_template(prompt_file) else: prompt_template = tool_agent_prompt if prompt_template == "": logger.error("No prompt template found") raise ValueError("No prompt template found") prompt = prompt_template.format( metadata=self.metadata, sample_records=sample_records, additional_instructions=additional_instructions ) response = run_prompt(prompt) content = response["choices"][0]["message"]["content"] # Remove markdown code blocks if present if content.startswith("```json"): content = content.replace("```json", "").replace("```", "").strip() elif content.startswith("```"): content = content.replace("```", "").strip() try: return json.loads(content) except json.JSONDecodeError as e: logger.error(f"Failed to parse JSON response: {e}") logger.error(f"Raw content: {content}") raise ValueError(f"Invalid JSON response from LLM: {e}") def save_tool_to_file(self, tool_spec: Dict[str, Any], file_path: str) -> None: """Save tool specification to a file.""" with open(file_path, 'w') as f: json.dump(tool_spec, f, indent=2) logger.info(f"Tool specification written to '{file_path}'") def main(): """Main function to parse arguments and generate tool specification.""" parser = argparse.ArgumentParser( description="Generate tool specification for Astra DB table", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: python generate_tool_spec.py --table-name users --db-name mydb --out-file user_tool.json python generate_tool_spec.py -t products -d ecommerce -o product_tool.json """ ) parser.add_argument( "--export-prompt", "-ep", required=False, help="Export the prompt used to generate the tool specification" ) parser.add_argument( "--table-name", "-t", required=False, help="Name of the table to analyze" ) parser.add_argument( "--collection-name", "-c", required=False, help="Name of the collection to analyze" ) parser.add_argument( "--keyspace-name", "-k", required=False, default="default_keyspace", help="Name of the keyspace to analyze" ) parser.add_argument( "--db-name", "-d", required=False, help="Name of the database containing the table" ) parser.add_argument( "--out-file", "-o", required=False, help="Output file path for the generated tool specification JSON" ) parser.add_argument( "--sample-size", "-s", type=int, default=5, help="Number of sample records to analyze (default: 5)" ) parser.add_argument( "--prompt-file", "-pf", required=False, help="Prompt file to use for the tool specification generation" ) parser.add_argument( "--additional-instructions", "-ai", required=False, help="Additional instructions for the tool specification generation" ) args = parser.parse_args() if args.export_prompt: export_prompt(args.export_prompt) sys.exit(0) if (args.table_name and args.collection_name): logger.error("Either --table-name or --collection-name must be provided") sys.exit(1) if (not args.out_file): logger.error("--out-file is required") sys.exit(1) token = os.getenv("ASTRA_DB_APPLICATION_TOKEN") db_name = os.getenv("ASTRA_DB_DB_NAME") tool_agent = AstraToolAgent(token, db_name, args.table_name, args.keyspace_name) tool_spec = tool_agent.generate_tool_specification(args.sample_size, args.additional_instructions, args.prompt_file) if args.out_file: tool_agent.save_tool_to_file(tool_spec, args.out_file) logger.info(f"Tool specification written to '{args.out_file}'") else: logger.info("Tool specification:") logger.info(json.dumps(tool_spec, indent=2)) sys.exit(1) if __name__ == "__main__": main()