rag_query_test.py

from langchain_openai import ChatOpenAI
from src.server.rag_query import RagQuery
from src.pgdb.knowledge.txt_doc_table import TxtDoc
from src.pgdb.knowledge.k_db import PostgresDB
from src.pgdb.knowledge.similarity import VectorStore_FAISS
import json
from src.config.consts import (
    EMBEEDING_MODEL_PATH,
    FAISS_STORE_PATH,
    INDEX_NAME,
    VEC_DB_HOST,
    VEC_DB_PASSWORD,
    VEC_DB_PORT,
    VEC_DB_USER,
    VEC_DB_DBNAME,
    SIMILARITY_SHOW_NUMBER,
)

base_llm = ChatOpenAI(
    openai_api_key='xxxxxxxxxxxxx',
    openai_api_base='http://192.168.10.14:8000/v1',
    model_name='Qwen2-7B',
    verbose=True
)

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)

k_db = PostgresDB(host=VEC_DB_HOST, database=VEC_DB_DBNAME, user=VEC_DB_USER, password=VEC_DB_PASSWORD, port=VEC_DB_PORT)
k_db.connect()
db = TxtDoc(k_db)

rag_query = RagQuery(base_llm=base_llm,_faiss_db=vecstore_faiss,_db=db)
question = "大通县和化隆县最大降雨量是多少"
history = ""
result = rag_query.query(question=question,history=history)
print(result)
j = json.loads(result["docs"], strict=False)
print(type(j))
print(j)