Spaces:
Runtime error
Runtime error
| from langchain.document_loaders.unstructured import UnstructuredFileLoader | |
| from langchain.text_splitter import CharacterTextSplitter | |
| from langchain.embeddings import OpenAIEmbeddings | |
| from langchain.vectorstores import Chroma | |
| from langchain.chains import RetrievalQA | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.schema import AIMessage, HumanMessage, SystemMessage, Document | |
| import os | |
| API_KEY = os.getenv("API_TOKEN") | |
| print("API:", API_KEY) | |
| class Agent: | |
| def __init__(self, args=None) -> None: | |
| self.embeddings = OpenAIEmbeddings(openai_api_key=API_KEY) | |
| self.llm = ChatOpenAI(temperature=0.5, openai_api_key=API_KEY) | |
| self.context_value = "" | |
| self.use_context = False | |
| def load_context(self, doc_path): | |
| loader = UnstructuredFileLoader(doc_path.name) | |
| print('Loading file:', doc_path.name) | |
| self.documents = loader.load() | |
| self.split() | |
| return f"Using file from {doc_path.name}" | |
| def split(self): | |
| text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0) | |
| texts = text_splitter.split_documents(self.documents) | |
| docsearch = Chroma.from_documents(texts, self.embeddings) | |
| self.qa = RetrievalQA.from_chain_type(llm=self.llm, chain_type="stuff", retriever=docsearch.as_retriever()) | |
| print("Context updated") | |
| self.use_context = True | |
| def asking(self, prompt): | |
| if self.use_context: | |
| print("Answering with your context") | |
| return self.qa.run(prompt) | |
| else: | |
| print("Answering without your context") | |
| return self.llm([HumanMessage(content=prompt)]).content | |
| def get_context(self, context): | |
| self.context_value = context | |
| self.documents = [Document(page_content=context, metadata={'source': ''})] | |
| self.split() | |
| self.use_context = True | |
| def load_context_file(self, file): | |
| print('Loading file:', file.name) | |
| text = '' | |
| for line in open(file.name, 'r', encoding='utf8'): | |
| text += line | |
| self.context_value = text | |
| return text | |
| def clear_context(self): | |
| self.context_value = "" | |
| self.use_context = False | |
| return "" | |