MyAIServ MCP Server

import os from datetime import datetime from typing import Any, Dict, List, Optional from elasticsearch import AsyncElasticsearch, NotFoundError class ElasticsearchStorage: """Класс для работы с Elasticsearch""" def __init__(self): self.es = AsyncElasticsearch( [os.getenv("ELASTICSEARCH_URL", "http://elasticsearch:9200")] ) self.indices = {"prompts": "mcp_prompts", "resources": "mcp_resources"} async def initialize(self): """Инициализация индексов""" # Маппинг для промптов prompt_mapping = { "mappings": { "properties": { "name": {"type": "keyword"}, "description": {"type": "text"}, "content": {"type": "text"}, "arguments": {"type": "nested"}, "vector": {"type": "dense_vector", "dims": 384}, "created_at": {"type": "date"}, "updated_at": {"type": "date"}, } } } # Маппинг для ресурсов resource_mapping = { "mappings": { "properties": { "uri": {"type": "keyword"}, "name": {"type": "text"}, "content": {"type": "text"}, "mime_type": {"type": "keyword"}, "vector": {"type": "dense_vector", "dims": 384}, "metadata": {"type": "object"}, } } } # Создаем индексы если их нет for index_name in self.indices.values(): if not await self.es.indices.exists(index=index_name): mapping = ( prompt_mapping if "prompts" in index_name else resource_mapping ) await self.es.indices.create(index=index_name, body=mapping) async def close(self): """Закрытие соединения""" await self.es.close() # Методы для работы с промптами async def index_prompt(self, prompt_data: Dict[str, Any]) -> str: """Индексация промпта""" prompt_data["created_at"] = datetime.utcnow() prompt_data["updated_at"] = datetime.utcnow() result = await self.es.index( index=self.indices["prompts"], document=prompt_data ) return result["_id"] async def get_prompt(self, prompt_id: str) -> Optional[Dict[str, Any]]: """Получение промпта по ID""" try: result = await self.es.get(index=self.indices["prompts"], id=prompt_id) return result["_source"] except NotFoundError: return None except Exception as e: print(f"Error getting prompt {prompt_id}: {e}") return None async def search_prompts(self, query: str, size: int = 10) -> List[Dict[str, Any]]: """Поиск промптов""" body = { "query": { "multi_match": { "query": query, "fields": ["name^2", "description", "content"], } }, "size": size, } result = await self.es.search(index=self.indices["prompts"], body=body) return [hit["_source"] for hit in result["hits"]["hits"]] # Методы для работы с ресурсами async def index_resource(self, resource_data: Dict[str, Any]) -> str: """Индексация ресурса""" result = await self.es.index( index=self.indices["resources"], document=resource_data ) return result["_id"] async def get_resource(self, resource_uri: str) -> Optional[Dict[str, Any]]: """Получение ресурса по URI""" try: result = await self.es.search( index=self.indices["resources"], body={"query": {"term": {"uri.keyword": resource_uri}}}, ) hits = result["hits"]["hits"] return hits[0]["_source"] if hits else None except Exception as e: print(f"Error getting resource {resource_uri}: {e}") return None async def search_resources( self, query: str, mime_type: Optional[str] = None, size: int = 10 ) -> List[Dict[str, Any]]: """Поиск ресурсов""" should_queries = [ {"multi_match": {"query": query, "fields": ["name^2", "content"]}} ] if mime_type: should_queries.append({"term": {"mime_type.keyword": mime_type}}) body = { "query": {"bool": {"should": should_queries, "minimum_should_match": 1}}, "size": size, } result = await self.es.search(index=self.indices["resources"], body=body) return [hit["_source"] for hit in result["hits"]["hits"]] async def delete_resource(self, resource_uri: str) -> bool: """Удаление ресурса""" try: await self.es.delete_by_query( index=self.indices["resources"], body={"query": {"term": {"uri.keyword": resource_uri}}}, ) return True except Exception as e: print(f"Error deleting resource {resource_uri}: {e}") return False # Создаем глобальный экземпляр es_storage = ElasticsearchStorage()