86 lines
3.8 KiB
Python
86 lines
3.8 KiB
Python
import logging
|
|
import time
|
|
from zhipuai import ZhipuAI
|
|
from app.utils.prompt_repository import PromptRepository
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ZhipuKbService:
|
|
def __init__(self):
|
|
self.model_name = "glm-4-plus"
|
|
self.app_secret_key = "d54f764a1d67c17d857bd3983b772016.GRjowY0fyiMNurLc"
|
|
logger.info("ZhipuKbService initialized with model: %s", self.model_name)
|
|
|
|
|
|
def retrive(self, message, knowledge_id, prompt_template):
|
|
logger.info("Starting retrive call with knowledge_id: %s", knowledge_id)
|
|
start_time = time.time()
|
|
client = ZhipuAI(api_key=self.app_secret_key)
|
|
default_prompt = "从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找问题\n\"\"\"\n{{question}}\n\"\"\"\n的答案，找到答案就仅使用文档语句回答问题，找不到答案就用自身知识回答并且告诉用户该信息不是来自文档。\n不要复述问题，直接开始回答。你是一个智能助理，请用该角色的语言风格和对话方式回答问题。"
|
|
|
|
if prompt_template is None or prompt_template == "":
|
|
prompt_template = default_prompt
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model=self.model_name,
|
|
messages=[
|
|
{"role": "user", "content": message},
|
|
],
|
|
tools=[
|
|
{
|
|
"type": "retrieval",
|
|
"retrieval": {
|
|
"knowledge_id": knowledge_id,
|
|
"prompt_template": prompt_template
|
|
}
|
|
}
|
|
],
|
|
stream=False,
|
|
max_tokens=4095,
|
|
temperature=0.01, # default=0.01
|
|
top_p=0.1, #default=0.1
|
|
)
|
|
result = response.choices[0].message.content
|
|
end_time = time.time()
|
|
logger.info("retrive call completed in %.2f seconds", end_time - start_time)
|
|
return result
|
|
except Exception as e:
|
|
logger.error("Error in retrive: %s", str(e))
|
|
raise
|
|
|
|
def retrive_sse(self, message, knowledge_id, prompt_template=None,system_prompt=None):
|
|
logger.info("Starting retrive_sse call with knowledge_id: %s, message:%s", knowledge_id, message)
|
|
start_time = time.time()
|
|
client = ZhipuAI(api_key=self.app_secret_key)
|
|
default_prompt = "从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找问题\n\"\"\"\n{{question}}\n\"\"\"\n的答案，找到答案就仅使用文档语句回答问题，找不到答案就告诉用户知识库中没有该信息。\n不要复述问题，直接开始回答。"
|
|
messages = [{"role": "user", "content": message}]
|
|
# if system_prompt != None:
|
|
# messages.append({"role": "system", "content": system_prompt})
|
|
if prompt_template is None or prompt_template == "":
|
|
prompt_template = default_prompt
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model=self.model_name,
|
|
messages=messages,
|
|
tools=[
|
|
{
|
|
"type": "retrieval",
|
|
"retrieval": {
|
|
"knowledge_id": knowledge_id,
|
|
"prompt_template": prompt_template
|
|
}
|
|
}
|
|
],
|
|
stream=True,
|
|
max_tokens=4095,
|
|
temperature=0.01, # default=0.01
|
|
top_p=0.1, #default=0.1
|
|
)
|
|
for chunk in response:
|
|
yield chunk.choices[0].delta.content
|
|
end_time = time.time()
|
|
logger.info("retrive_sse call completed in %.2f seconds", end_time - start_time)
|
|
except Exception as e:
|
|
logger.error("Error in retrive_sse: %s", str(e))
|
|
raise
|