|
|
|
|
|
""" |
|
|
Created on Sun Feb 2 20:31:22 2025 |
|
|
|
|
|
@author: BM109X32G-10GPU-02 |
|
|
""" |
|
|
|
|
|
|
|
|
from langchain.tools import BaseTool |
|
|
import os |
|
|
from langchain.prompts.chat import ( |
|
|
ChatPromptTemplate, |
|
|
HumanMessagePromptTemplate, |
|
|
SystemMessagePromptTemplate, |
|
|
) |
|
|
|
|
|
from langchain import PromptTemplate |
|
|
from langchain import HuggingFacePipeline |
|
|
|
|
|
from langchain.base_language import BaseLanguageModel |
|
|
|
|
|
|
|
|
|
|
|
from langchain.chains import RetrievalQA |
|
|
|
|
|
from langchain_community.document_loaders import PyPDFLoader |
|
|
|
|
|
from langchain_openai import ChatOpenAI |
|
|
from langchain_community.vectorstores import FAISS |
|
|
from torch import cuda, bfloat16 |
|
|
device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu' |
|
|
from langchain_openai import OpenAIEmbeddings |
|
|
|
|
|
|
|
|
class rag(BaseTool): |
|
|
name: str = "RAG" |
|
|
description: str= ( |
|
|
"Useful to answer questions that require technical " |
|
|
|
|
|
"Provide specialized knowledge information for solving Q&A questions" |
|
|
"Input query , return the response" |
|
|
|
|
|
) |
|
|
|
|
|
llm: BaseLanguageModel = None |
|
|
path : str = None |
|
|
|
|
|
def __init__(self, path: str = None): |
|
|
super().__init__( ) |
|
|
self.llm = ChatOpenAI(model="gpt-4o-2024-11-20",api_key=os.getenv("OPENAI_API_KEY"), |
|
|
base_url=os.getenv("OPENAI_API_BASE") |
|
|
) |
|
|
self.path = path |
|
|
|
|
|
|
|
|
def _run(self, query ) -> str: |
|
|
embeddings = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"), |
|
|
base_url=os.getenv("OPENAI_API_BASE")) |
|
|
vectorstore=FAISS.load_local(r"rag", embeddings,allow_dangerous_deserialization =True) |
|
|
|
|
|
|
|
|
template = """ |
|
|
|
|
|
You are an expert chemist and your task is to respond to the question or |
|
|
solve the problem to the best of your ability.You can only respond with a single "Final Answer" format. |
|
|
You need to list the key points and explain them in detail and accurately |
|
|
Use the following pieces of context to answer the question at the end. |
|
|
If you don't know the answer, just say that you don't know, don't try to make up an answer. |
|
|
<context> |
|
|
{context} |
|
|
</context> |
|
|
|
|
|
Question: {question} |
|
|
Answer: |
|
|
|
|
|
""" |
|
|
|
|
|
prompt = PromptTemplate(template=template, input_variables=[ "question"]) |
|
|
qa_chain = RetrievalQA.from_chain_type( |
|
|
llm=self.llm, |
|
|
chain_type="stuff", |
|
|
retriever=vectorstore.as_retriever(search_kwargs={"k": 5}), |
|
|
return_source_documents=False, |
|
|
|
|
|
chain_type_kwargs={"prompt": prompt}, |
|
|
) |
|
|
chat_history = [] |
|
|
|
|
|
|
|
|
result = qa_chain.invoke(query) |
|
|
return result['result'] |
|
|
|
|
|
|
|
|
async def _arun(self, query) -> str: |
|
|
"""Use the tool asynchronously.""" |
|
|
raise NotImplementedError("this tool does not support async") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
