75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
|
from app.service.lookbooks.config.config import DOCUMENT_COLLECTION # 引入向量数据库
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
async def query_lookbooks_service(tag: Optional[str], year: Optional[str], n_results: int):
|
|||
|
|
"""
|
|||
|
|
查询向量数据库，支持按 tag 和 year 过滤
|
|||
|
|
:param tag: 查询过滤的标签
|
|||
|
|
:param year: 查询过滤的年份
|
|||
|
|
:param n_results: 返回的结果数量
|
|||
|
|
:return: 查询结果列表
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
# 选择一个主要过滤条件进行初步查询
|
|||
|
|
primary_filter = {}
|
|||
|
|
if tag:
|
|||
|
|
primary_filter['tag'] = tag
|
|||
|
|
elif year:
|
|||
|
|
primary_filter['year'] = year
|
|||
|
|
else:
|
|||
|
|
primary_filter = None
|
|||
|
|
|
|||
|
|
# 如果没有任何条件，直接返回空列表
|
|||
|
|
if not primary_filter:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 使用主过滤条件进行查询
|
|||
|
|
query_result = DOCUMENT_COLLECTION.get(where=primary_filter)
|
|||
|
|
|
|||
|
|
# 打印 query_result 以检查数据结构是否符合预期
|
|||
|
|
print(query_result) # 或者使用 logger 记录以调试
|
|||
|
|
|
|||
|
|
# 检查 query_result 的结构
|
|||
|
|
if not isinstance(query_result, dict) or 'documents' not in query_result:
|
|||
|
|
raise ValueError("Expected query_result to be a dict containing a 'documents' key")
|
|||
|
|
|
|||
|
|
documents = query_result['documents']
|
|||
|
|
|
|||
|
|
# 确保 documents 是一个列表
|
|||
|
|
if not isinstance(documents, list):
|
|||
|
|
raise ValueError("Expected 'documents' to be a list")
|
|||
|
|
|
|||
|
|
# 对文档进行进一步过滤
|
|||
|
|
filtered_results = []
|
|||
|
|
for item in documents:
|
|||
|
|
# 检查每个 item 是否是字典，如果不是字典，可能是简单文本描述
|
|||
|
|
if isinstance(item, dict):
|
|||
|
|
metadata = item.get('metadata', {})
|
|||
|
|
if (not year or metadata.get('year') == year) and (not tag or metadata.get('tag') == tag):
|
|||
|
|
filtered_results.append(item)
|
|||
|
|
elif isinstance(item, str):
|
|||
|
|
# 如果 item 是字符串，假设它是描述文本，构造默认的 metadata
|
|||
|
|
filtered_results.append({
|
|||
|
|
"text": item,
|
|||
|
|
"metadata": {}
|
|||
|
|
})
|
|||
|
|
else:
|
|||
|
|
# 如果是其他数据类型，打印警告并跳过
|
|||
|
|
print(f"Unexpected item type in documents: {type(item)}")
|
|||
|
|
|
|||
|
|
# 限制结果数量
|
|||
|
|
limited_results = filtered_results[:n_results]
|
|||
|
|
|
|||
|
|
# 格式化结果
|
|||
|
|
result_list = []
|
|||
|
|
for item in limited_results:
|
|||
|
|
result_list.append({
|
|||
|
|
"description": item.get('text', ""),
|
|||
|
|
"metadata": item.get('metadata', {})
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return result_list
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
raise e # 抛出异常，让接口层处理
|