mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-10-21 17:09:46 +00:00
30 lines
987 B
Python
30 lines
987 B
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from pilot.vector_store.file_loader import KnownLedge2Vector
|
|
from langchain.prompts import PromptTemplate
|
|
from pilot.conversation import conv_qa_prompt_template
|
|
from pilot.configs.model_config import VECTOR_SEARCH_TOP_K
|
|
from pilot.model.vicuna_llm import VicunaLLM
|
|
|
|
class KnownLedgeBaseQA:
|
|
|
|
def __init__(self) -> None:
|
|
k2v = KnownLedge2Vector()
|
|
self.vector_store = k2v.init_vector_store()
|
|
self.llm = VicunaLLM()
|
|
|
|
def get_similar_answer(self, query):
|
|
|
|
prompt = PromptTemplate(
|
|
template=conv_qa_prompt_template,
|
|
input_variables=["context", "question"]
|
|
)
|
|
|
|
retriever = self.vector_store.as_retriever(search_kwargs={"k": VECTOR_SEARCH_TOP_K})
|
|
docs = retriever.get_relevant_documents(query=query)
|
|
|
|
context = [d.page_content for d in docs]
|
|
result = prompt.format(context="\n".join(context), question=query)
|
|
return result
|