File size: 2,245 Bytes
e011405
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ""