Spaces:
Runtime error
Runtime error
File size: 2,245 Bytes
56523b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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 ""
|