from .utils import catch_error_and_return
from dotenv import load_dotenv
from fastmcp import FastMCP
import requests
import json
import os
load_dotenv()
llm_mcp = FastMCP(name="LLM server")
@llm_mcp.tool
@catch_error_and_return
def llm_chat_completion(text: str) -> str:
"""
Generate answer for text
:param text: text
:return: answer
"""
url = os.getenv("LLM_SERVICE_CHAT_COMPLETIONS_URL")
api_key = os.getenv("LLM_SERVICE_API_KEY")
model_name = os.getenv("LLM_SERVICE_MODEL")
data = {"model": model_name, "messages": [{"content": text, "role": "user"}]}
headers = {
"Authorization": f"Bearer {api_key}",
}
result = requests.post(
url,
data=json.dumps(data),
headers=headers,
).json()
return result["choices"][0]["message"]["content"]
@llm_mcp.tool
@catch_error_and_return
def llm_get_models() -> dict:
"""
Get existing models for generate answer
:return: existing models for generate answer
"""
url = os.getenv("LLM_SERVICE_MODELS_URL")
api_key = os.getenv("LLM_SERVICE_API_KEY")
headers = {
"Authorization": f"Bearer {api_key}",
}
result = requests.get(url, headers=headers).json()
return result
@llm_mcp.tool
@catch_error_and_return
def llm_stream_completion(text: str) -> str:
"""
Streaming generate text
:param text: text for LLM
:return:
"""
url = os.getenv("LLM_SERVICE_COMPLETIONS_URL")
api_key = os.getenv("LLM_SERVICE_API_KEY")
model = os.getenv("LLM_SERVICE_MODEL")
data = {"model": model, "prompt": [text], "stream": True}
headers = {
"Authorization": f"Bearer {api_key}",
}
# TODO: check stream correct?
response = requests.post(url, headers=headers, json=data, stream=True)
answer = ""
for line in response.iter_lines():
if line:
decoded_line = line.decode("utf-8")[6:]
if decoded_line == "[DONE]":
break
answer += json.loads(decoded_line)["choices"][0]["text"]
return answer
@llm_mcp.tool
def health_check():
"""
health check for server
:return: server status
"""
url = os.getenv("LLM_SERVICE_HEALTH_URL")
api_key = os.getenv("LLM_SERVICE_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}"}
