from .utils import catch_error_and_return
from dotenv import load_dotenv
from fastmcp import FastMCP
from typing import List
import requests
import logging
import json
import os
load_dotenv()
logger = logging.getLogger(__name__)
embedding_mcp = FastMCP(name="Embedding server")
@embedding_mcp.tool
@catch_error_and_return
def embedding_generate(text: str) -> dict:
"""
Method for gets embeddings text
:param text: text
:return: List of lists embeddings
"""
url = os.getenv("EMBEDDING_URL")
api_key = os.getenv("EMBEDDING_API_KEY")
model_name = os.getenv("EMBEDDING_MODEL_NAME")
data = {"model": model_name, "input": [text]}
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
result = requests.post(
url,
data=json.dumps(data),
headers=headers,
).json()
return result
@embedding_mcp.tool
@catch_error_and_return
def embedding_batch_generate(texts: List[str]) -> dict:
"""
Method for gets embeddings text
:param texts: List of texts
:return: List of lists embeddings
"""
url = os.getenv("EMBEDDING_URL")
api_key = os.getenv("EMBEDDING_API_KEY")
model_name = os.getenv("EMBEDDING_MODEL_NAME")
data = {"model": model_name, "input": texts}
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
result = requests.post(
url,
data=json.dumps(data),
headers=headers,
).json()
return result
@embedding_mcp.tool
@catch_error_and_return
def embedding_get_models() -> dict:
"""
Method for get all embeddings models
:return: existing embedding models
"""
url = os.getenv("EMBEDDING_URL_MODELS")
api_key = os.getenv("EMBEDDING_API_KEY")
headers = {
"Authorization": f"Bearer {api_key}",
}
result = requests.get(url, headers=headers).json()
return result
@embedding_mcp.tool
def health_check():
"""
health check for server
:return: server status
"""
url = os.getenv("EMBEDDING_HEALTH_URL")
api_key = os.getenv("EMBEDDING_API_KEY")
headers = {
"Authorization": f"Bearer {api_key}",
}
try:
result = requests.get(url, headers=headers)
result.raise_for_status()
return {"status": "ok"}
except Exception as ex:
return {"status": f"Error {ex}"}
if __name__ == "__main__":
embedding_mcp.run(transport="stdio")
