|
|
|
@ -3,6 +3,7 @@ from psycopg2.extras import execute_values
|
|
|
|
import numpy as np
|
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
|
|
|
from env import PgvectorDataBaseConfig
|
|
|
|
from env import PgvectorDataBaseConfig
|
|
|
|
|
|
|
|
from util.log_util import logger
|
|
|
|
|
|
|
|
|
|
|
|
# 数据库连接配置
|
|
|
|
# 数据库连接配置
|
|
|
|
DB_CONFIG = {
|
|
|
|
DB_CONFIG = {
|
|
|
|
@ -18,10 +19,10 @@ def connect_to_db():
|
|
|
|
"""建立数据库连接"""
|
|
|
|
"""建立数据库连接"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
|
|
print("成功连接到数据库")
|
|
|
|
logger.info("成功连接到数据库")
|
|
|
|
return conn
|
|
|
|
return conn
|
|
|
|
except Exception as e:
|
|
|
|
except Exception as e:
|
|
|
|
print(f"连接数据库失败: {e}")
|
|
|
|
logger.error(f"连接数据库失败: {e}")
|
|
|
|
return None
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -51,9 +52,9 @@ def setup_vector_extension(conn):
|
|
|
|
);
|
|
|
|
);
|
|
|
|
""")
|
|
|
|
""")
|
|
|
|
conn.commit()
|
|
|
|
conn.commit()
|
|
|
|
print("成功设置向量扩展和表结构")
|
|
|
|
# print("成功设置向量扩展和表结构")
|
|
|
|
except Exception as e:
|
|
|
|
except Exception as e:
|
|
|
|
print(f"设置数据库结构失败: {e}")
|
|
|
|
logger.error(f"设置数据库结构失败: {e}")
|
|
|
|
conn.rollback()
|
|
|
|
conn.rollback()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -68,9 +69,9 @@ def insert_vectors(conn, data):
|
|
|
|
data
|
|
|
|
data
|
|
|
|
)
|
|
|
|
)
|
|
|
|
conn.commit()
|
|
|
|
conn.commit()
|
|
|
|
print(f"成功插入 {len(data)} 条向量数据")
|
|
|
|
logger.info(f"成功插入 {len(data)} 条向量数据")
|
|
|
|
except Exception as e:
|
|
|
|
except Exception as e:
|
|
|
|
print(f"插入数据失败: {e}")
|
|
|
|
logger.error(f"插入数据失败: {e}")
|
|
|
|
conn.rollback()
|
|
|
|
conn.rollback()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -88,7 +89,7 @@ def search_similar_table(conn, query_vector, limit=5):
|
|
|
|
results = cur.fetchall()
|
|
|
|
results = cur.fetchall()
|
|
|
|
return results
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
|
|
except Exception as e:
|
|
|
|
print(f"搜索相似向量失败: {e}")
|
|
|
|
logger.error(f"搜索相似向量失败: {e}")
|
|
|
|
return []
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
|
|
|
def search_similar_case(conn, query_vector, limit=5):
|
|
|
|
def search_similar_case(conn, query_vector, limit=5):
|
|
|
|
@ -105,7 +106,7 @@ def search_similar_case(conn, query_vector, limit=5):
|
|
|
|
results = cur.fetchall()
|
|
|
|
results = cur.fetchall()
|
|
|
|
return results
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
|
|
except Exception as e:
|
|
|
|
print(f"搜索相似向量失败: {e}")
|
|
|
|
logger.error(f"搜索相似向量失败: {e}")
|
|
|
|
return []
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
|
|
|
def main():
|
|
|
|
def main():
|
|
|
|
@ -133,9 +134,9 @@ def main():
|
|
|
|
query_vector = [1.0, 1.0, 1.0]
|
|
|
|
query_vector = [1.0, 1.0, 1.0]
|
|
|
|
similar_docs = search_similar_table(conn, query_vector)
|
|
|
|
similar_docs = search_similar_table(conn, query_vector)
|
|
|
|
|
|
|
|
|
|
|
|
print("\n相似文档搜索结果:")
|
|
|
|
logger.info("\n相似文档搜索结果:")
|
|
|
|
for content, distance in similar_docs:
|
|
|
|
for content, distance in similar_docs:
|
|
|
|
print(f"内容: {content}, 距离: {distance}")
|
|
|
|
logger.info(f"内容: {content}, 距离: {distance}")
|
|
|
|
|
|
|
|
|
|
|
|
# 关闭连接
|
|
|
|
# 关闭连接
|
|
|
|
conn.close()
|
|
|
|
conn.close()
|
|
|
|
|