122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
import logging
|
|
import requests
|
|
from requests.auth import HTTPBasicAuth
|
|
from urllib.parse import parse_qs
|
|
from typing import Optional
|
|
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain_core.runnables import RunnablePassthrough
|
|
|
|
from chatchat.server.pydantic_v1 import Field
|
|
from chatchat.server.utils import get_tool_config, get_ChatOpenAI
|
|
# from chatchat.server.api_server.chat_routes import global_model_name
|
|
# from chatchat.configs import (
|
|
# MAX_TOKENS,
|
|
# )
|
|
|
|
from .tools_registry import BaseToolOutput, regist_tool
|
|
|
|
logger = logging.getLogger()
|
|
|
|
# Prompt for the prom_chain
|
|
PROMETHEUS_PROMPT_TEMPLATE = """ You are an expert in Prometheus, a powerful time-series monitoring service.
|
|
Your main task is to translate users' specific requirements into PromQL queries.
|
|
This includes understanding their monitoring needs, the specific metrics they are interested in,
|
|
the time period for which they want the data, and any specific conditions or thresholds they want to apply.
|
|
Your goal is to provide the most accurate and efficient PromQL query based on the given information.
|
|
|
|
Please return the PromQL in a format that can be used with the HTTP API, such as:
|
|
|
|
'query?query=up&time=2015-07-01T20:10:51.781Z' for instant data queries.
|
|
'query_range?query=up&start=2015-07-01T20:10:30.781Z&end=2015-07-01T20:11:00.781Z&step=15s' for range data queries.
|
|
|
|
I will automatically fill in the Prometheus IP and port. Please provide the query according to the example,
|
|
and no other content is needed.
|
|
|
|
Question: {query}
|
|
"""
|
|
|
|
|
|
def execute_promql_request(url: str, params: dict, auth: Optional[HTTPBasicAuth], promql: str) -> str:
|
|
try:
|
|
response = requests.get(url, params=params, auth=auth)
|
|
except requests.exceptions.RequestException as e:
|
|
return f"PromQL: {promql} 的错误: {e}\n"
|
|
|
|
if response.status_code == 200:
|
|
return f"PromQL: {promql} 的查询结果: {response.json()}\n"
|
|
else:
|
|
return f"PromQL: {promql} 的错误: {response.text}\n"
|
|
|
|
|
|
def query_prometheus(query: str, config: dict) -> str:
|
|
prometheus_endpoint = config["prometheus_endpoint"]
|
|
username = config["username"]
|
|
password = config["password"]
|
|
|
|
llm = get_ChatOpenAI(
|
|
# model_name=global_model_name,
|
|
# temperature=0.1,
|
|
# streaming=True,
|
|
# local_wrap=True,
|
|
# verbose=True,
|
|
# max_tokens=MAX_TOKENS,
|
|
)
|
|
|
|
prometheus_prompt = ChatPromptTemplate.from_template(PROMETHEUS_PROMPT_TEMPLATE)
|
|
output_parser = StrOutputParser()
|
|
|
|
prometheus_chain = (
|
|
{"query": RunnablePassthrough()}
|
|
| prometheus_prompt
|
|
| llm
|
|
| output_parser
|
|
)
|
|
|
|
# 根据用户名和密码是否存在来设置 auth 参数
|
|
auth = HTTPBasicAuth(username, password) if username and password else None
|
|
|
|
promql = prometheus_chain.invoke(query)
|
|
logger.info(f"PromQL: {promql}")
|
|
|
|
# debug
|
|
# promql = 'query?query=kube_pod_status_phase{namespace="default"}'
|
|
|
|
# 从 promql 字符串中提取 query_type 和 参数
|
|
try:
|
|
query_type, query_params = promql.split('?')
|
|
except ValueError as e:
|
|
logger.error(f"Promql split error: {e}")
|
|
content = f"PromQL: {promql} 的错误: {e}\n"
|
|
return content
|
|
|
|
prometheus_url = f"{prometheus_endpoint}/api/v1/{query_type}"
|
|
|
|
# params 的格式如下:
|
|
"""
|
|
{
|
|
"query": "kube_pod_status_phase{namespace=\"default\"}"
|
|
}
|
|
"""
|
|
params = {k: v[0] for k, v in parse_qs(query_params).items()}
|
|
|
|
# 向 Prometheus 发起 HTTP 请求
|
|
content = execute_promql_request(prometheus_url, params, auth, promql)
|
|
|
|
logger.info(content)
|
|
return content
|
|
|
|
|
|
# @regist_tool(title="Prometheus对话")
|
|
def text2promql(
|
|
query: str = Field(
|
|
description="Tool for querying a Prometheus server, No need for PromQL statements, "
|
|
"just input the natural language that you want to chat with prometheus"
|
|
)
|
|
):
|
|
"""Use this tool to chat with prometheus, Input natural language,
|
|
then it will convert it into PromQL and execute it in the prometheus, then return the execution result."""
|
|
tool_config = get_tool_config("text2promql")
|
|
return BaseToolOutput(query_prometheus(query=query, config=tool_config))
|