Merge remote-tracking branch 'origin/geo' into geo

test/k_store_test.py

@@ -17,7 +17,7 @@ from src.config.consts import (
 from src.loader.callback import BaseCallback
 from src.server.rerank import BgeRerank
 
-sys.path.append("../")
+import time
 
 
 # 当返回值中带有“思考题”字样的时候，默认将其忽略。
@@ -81,6 +81,19 @@ def test_faiss_load():
         reset=False)
     print(vecstore_faiss.join_document(vecstore_faiss.get_text_similarity("我国什么时候全面开放低空领域")))
 
+def test_faiss_search():
+    vecstore_faiss = VectorStore_FAISS(
+        embedding_model_name=EMBEEDING_MODEL_PATH,
+        store_path=FAISS_STORE_PATH,
+        index_name=INDEX_NAME,
+        info={"port": VEC_DB_PORT, "host": VEC_DB_HOST, "dbname": VEC_DB_DBNAME, "username": VEC_DB_USER,
+              "password": VEC_DB_PASSWORD},
+        show_number=SIMILARITY_SHOW_NUMBER,
+        reset=False)
+    tic = time.perf_counter()
+    print(vecstore_faiss.get_text_similarity("化隆县的降雨情况如何"))
+    toc = time.perf_counter() - tic
+    print("Time:", toc)
 
 def test_reranker():
     from langchain_core.documents import Document
@@ -93,4 +106,5 @@ def test_reranker():
 if __name__ == "__main__":
     # test_faiss_from_dir()
     # test_faiss_load()
-    test_reranker()
+    # test_reranker()
+    test_faiss_search()

test/rag_test.py

@@ -13,13 +13,18 @@ base_llm = ChatOpenAI(
     verbose=True
 )
 
-ext = QAExt(base_llm)
-question = "明天适合去吗？"
-message = [
+def test_qaext():
+    ext = QAExt(base_llm)
+    question = "明天适合去吗？"
+    message = [
         ("我们明天去爬山吧", "好呀"),
         ("天气怎么样", "天气晴朗"),
-]
+    ]
 
     
-result =  ext.extend_query(question, message)
-print(result.content)
+    result =  ext.extend_query(question, message)
+    print(result.content)
+
+
+if __name__ == "__main__":
+    pass
\ No newline at end of file