212 lines
8.4 KiB
Python
212 lines
8.4 KiB
Python
import urllib
|
||
|
||
from langchain_core.pydantic_v1 import BaseModel, Extra, Field, root_validator
|
||
import requests
|
||
from typing import Dict, List
|
||
from configs import kgo_search_url, kgo_professional_search_url
|
||
from configs import LLM_MODELS, TEMPERATURE, MAX_TOKENS
|
||
from server.chat.policy_fun_iast import get_llm_model_response
|
||
from datetime import datetime
|
||
import json
|
||
from configs.basic_config import *
|
||
class KgoSearchAPIWrapper(BaseModel):
|
||
search_kwargs: dict = Field(default_factory=dict)
|
||
|
||
search_map = {
|
||
"全部": "1000",
|
||
"期刊论文": "1002",
|
||
"学位论文": "1003",
|
||
"会议论文": "1004",
|
||
"政策": "1005",
|
||
"成果": "1006",
|
||
"科技成果": "1006",
|
||
"项目": "1007",
|
||
"报告": "1010",
|
||
"图书": "1011",
|
||
"外文期刊论文": "1013",
|
||
"外文资料": "1013",
|
||
"期刊": "1015",
|
||
"专利": "3001",
|
||
"新闻": "4001",
|
||
"数据集": "4004",
|
||
"视频": "4005",
|
||
"统计数据": "6001"
|
||
}
|
||
|
||
class Config:
|
||
"""Configuration for this pydantic object."""
|
||
|
||
extra = Extra.forbid
|
||
|
||
def _kgo_search_results(self, kgo_search_type: str, search_text: str) -> List[dict]:
|
||
params = {
|
||
"type": kgo_search_type,
|
||
"model": 20,
|
||
"text": search_text,
|
||
**self.search_kwargs,
|
||
}
|
||
response = requests.get(
|
||
kgo_search_url,
|
||
params=params, # type: ignore
|
||
)
|
||
response.raise_for_status()
|
||
search_results = response.json()
|
||
# print(search_results)
|
||
if "code" in search_results and search_results["code"] == 10000:
|
||
return search_results["results"]["datas"]
|
||
return []
|
||
|
||
def _kgo_professional_search_results(self, kgo_search_type: str, search_text: str, origin_query:str) -> List[dict]:
|
||
# 将逗号排除在编码范围之外，不进行编码
|
||
# encoded_kgo_search_type = ','.join(kgo_search_type.split(','))
|
||
# params = {
|
||
# "model": 20,
|
||
# "express": "theme=" + search_text,
|
||
# **self.search_kwargs,
|
||
# }
|
||
year = get_llm_model_response(
|
||
strategy_name="kgo_get_year",
|
||
llm_model_name=LLM_MODELS[0],
|
||
template_prompt_name="kgo_get_year",
|
||
prompt_param_dict={
|
||
"query": origin_query,
|
||
"year": datetime.now().strftime("%Y")
|
||
},
|
||
temperature=TEMPERATURE,
|
||
max_tokens=MAX_TOKENS
|
||
)
|
||
logger.info(f'KGO资源发布时间表达式：{year}')
|
||
# res = {}
|
||
# for response in solve_problem(query=search_text, conversation_id="", history=[], model_name=LLM_MODELS[0],temperature=TEMPERATURE,max_tokens=MAX_TOKENS,prompt_name="kgo_get_year",stream=False):
|
||
# res["text"] += json.loads(response)["text"]
|
||
params = {
|
||
"model": 20,
|
||
"express": "theme=" + search_text + " AND " + year,
|
||
**self.search_kwargs,
|
||
}
|
||
# 使用 urllib.parse.urlencode 构建查询参数部分
|
||
query_params = urllib.parse.urlencode(params)
|
||
# 构建完整的 URL
|
||
url = f"{kgo_professional_search_url}?type={kgo_search_type}&{query_params}"
|
||
logger.info(f'KGO专业检索请求URL: {url}')
|
||
try:
|
||
# 发起网络请求
|
||
response = requests.get(url)
|
||
response.raise_for_status()
|
||
search_results = response.json()
|
||
# 检查返回结果
|
||
if "code" in search_results and search_results["code"] == 10000:
|
||
if "datas" in search_results["results"]:
|
||
return search_results["results"]["datas"]
|
||
except requests.RequestException as e:
|
||
# 处理请求异常
|
||
print(f"请求异常: {e}")
|
||
|
||
# 如果出现任何错误或未找到结果，返回空列表
|
||
return []
|
||
|
||
@root_validator(pre=True)
|
||
def validate_environment(cls, values: Dict) -> Dict:
|
||
"""Validate that python package exists in environment."""
|
||
return values
|
||
|
||
def run(self, text: str, kgo_search_type: str) -> str:
|
||
"""Run query through Kgo Search and parse result."""
|
||
snippets = []
|
||
results = self._kgo_search_results(search_text=text, kgo_search_type=kgo_search_type)
|
||
if len(results) == 0:
|
||
return "No good Kgo Search Result was found"
|
||
for result in results:
|
||
snippets.append(result["snippet"])
|
||
|
||
return " ".join(snippets)
|
||
|
||
def results(self, text: str, kgo_search_type: str,origin_query:str) -> List[Dict]:
|
||
"""Run query through Kgo Search and return metadata.
|
||
|
||
Args:
|
||
text: The query to search for.
|
||
kgo_search_type: The type of results in search_map to return.
|
||
|
||
Returns:
|
||
A list of dictionaries with the following keys:
|
||
snippet - The description of the result.
|
||
title - The title of the result.
|
||
link - The link to the result.
|
||
author - The author of the result.
|
||
"""
|
||
metadata_results = []
|
||
results = self._kgo_professional_search_results(search_text=text, kgo_search_type=kgo_search_type,origin_query=origin_query)
|
||
|
||
if not results:
|
||
return [{"Result": "No good Kgo Search Result was found"}]
|
||
i = 0
|
||
for result in results:
|
||
i += 1
|
||
snippet = result.get("AB", "")
|
||
title = result.get("title", "")
|
||
publish_year = result.get("publish_year", "")
|
||
source_table_name = result.get("source_table_name", "")
|
||
resource_type = result.get("resource_type", "")
|
||
_id = result.get("_id", "")
|
||
author = ",".join(result.get("author", []))
|
||
keywords = ",".join(result.get("keywords", []))
|
||
|
||
link = ""
|
||
if source_table_name and resource_type:
|
||
link = f"https://kgo.ckcest.cn/kgo/detail/{self.search_map.get(resource_type, '')}/{source_table_name}/{_id}.html"
|
||
|
||
metadata_result = {
|
||
"snippet": snippet,
|
||
"title": title,
|
||
"link": link,
|
||
"author": author,
|
||
"keywords": keywords,
|
||
"publish_year": publish_year,
|
||
"resource_type": resource_type
|
||
}
|
||
metadata_results.append(metadata_result)
|
||
if i == 5:
|
||
break
|
||
# print("metadata_results:", metadata_results)
|
||
return metadata_results
|
||
|
||
def results_normal(self, text: str, kgo_search_type: str) -> List[Dict]:
|
||
metadata_results = []
|
||
results = self._kgo_search_results(search_text=text, kgo_search_type=kgo_search_type)
|
||
if len(results) == 0:
|
||
return [{"Result": "No good Kgo Search Result was found"}]
|
||
for result in results:
|
||
metadata_result = {
|
||
"snippet": result["AB"] if "AB" in result.keys() else "",
|
||
"publish_year": result["publish_year"] if "publish_year" in result.keys() else "",
|
||
"resource_type": result["resource_type"] if "resource_type" in result.keys() else "",
|
||
"title": result["title"] if "title" in result.keys() else "",
|
||
"link": "https://kgo.ckcest.cn/kgo/detail/" + self.search_map[result["resource_type"]]
|
||
+ "/" + result["source_table_name"] + "/" + result["_id"] + ".html"
|
||
if "source_table_name" in result.keys() and "resource_type" in result.keys() else "",
|
||
"keywords": result["keywords"] if "keywords" in result.keys() else "",
|
||
# "author": result["author_json"]["name"]
|
||
}
|
||
# print("metadata_result:", metadata_result)
|
||
metadata_results.append(metadata_result)
|
||
|
||
return metadata_results
|
||
|
||
|
||
if __name__ == '__main__':
|
||
search = KgoSearchAPIWrapper()
|
||
search_results = search.results("物理气相沉积", "1000")
|
||
print(search_results)
|
||
# docs = []
|
||
# for result in search_results:
|
||
# print(result["snippet"])
|
||
# # print("snippet:" + result["snippet"])
|
||
# # print("link:" + result["link"])
|
||
# # print("title:" + result["title"])
|
||
# doc = Document(page_content=result["snippet"] if "snippet" in result.keys() else "",
|
||
# metadata={"source": result["link"] if "link" in result.keys() else "",
|
||
# "filename": result["title"] if "title" in result.keys() else ""})
|
||
# docs.append(doc)
|
||
# print(docs)
|