Update app.py

app.py

@@ -1,36 +1,30 @@
 import os
 from typing import List
+from dotenv import load_dotenv
+
 from chainlit.types import AskFileResponse
 from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFLoader
 from aimakerspace.openai_utils.prompts import (
     UserRolePrompt,
     SystemRolePrompt,
-    AssistantRolePrompt,
 )
-from aimakerspace.openai_utils.embedding import EmbeddingModel
+
+from aimakerspace.embedding import EmbeddingModel
 from aimakerspace.vectordatabase import VectorDatabase
-from aimakerspace.openai_utils.chatmodel import ChatOpenAI
+from huggingface_utils.chatmodel import HuggingFaceLLM
 import chainlit as cl
 
-# Ensure OpenAI API key is set in environment variables
-if not os.getenv("OPENAI_API_KEY"):
-    raise ValueError("OPENAI_API_KEY environment variable is not set")
+# Загрузка токена из .env
+load_dotenv()
 
-system_template = """\
-Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
+system_template = """Use the following context to answer a user's question. If you cannot find the answer in the context, say you don't know."""
 system_role_prompt = SystemRolePrompt(system_template)
 
-user_prompt_template = """\
-Context:
-{context}
-
-Question:
-{question}
-"""
+user_prompt_template = """Context:\n{context}\n\nQuestion:\n{question}"""
 user_role_prompt = UserRolePrompt(user_prompt_template)
 
 class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
+    def __init__(self, llm: HuggingFaceLLM, vector_db_retriever: VectorDatabase) -> None:
         self.llm = llm
         self.vector_db_retriever = vector_db_retriever
 
@@ -42,7 +36,6 @@ class RetrievalAugmentedQAPipeline:
             context_prompt += context[0] + "\n"
 
         formatted_system_prompt = system_role_prompt.create_message()
-
         formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)
 
         async def generate_response():
@@ -53,45 +46,36 @@ class RetrievalAugmentedQAPipeline:
 
 text_splitter = CharacterTextSplitter()
 
-
 def process_file(file: AskFileResponse):
     import tempfile
     import shutil
-    
+
     print(f"Processing file: {file.name}")
-    
-    # Create a temporary file with the correct extension
+
     suffix = f".{file.name.split('.')[-1]}"
     with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
-        # Copy the uploaded file content to the temporary file
         shutil.copyfile(file.path, temp_file.name)
         print(f"Created temporary file at: {temp_file.name}")
-        
-        # Create appropriate loader
+
         if file.name.lower().endswith('.pdf'):
             loader = PDFLoader(temp_file.name)
         else:
             loader = TextFileLoader(temp_file.name)
-            
+
         try:
-            # Load and process the documents
             documents = loader.load_documents()
             texts = text_splitter.split_texts(documents)
             return texts
         finally:
-            # Clean up the temporary file
             try:
                 os.unlink(temp_file.name)
             except Exception as e:
                 print(f"Error cleaning up temporary file: {e}")
 
-
 @cl.on_chat_start
 async def on_chat_start():
     files = None
-
-    # Wait for the user to upload a file
-    while files == None:
+    while files is None:
         files = await cl.AskFileMessage(
             content="Please upload a Text or PDF file to begin!",
             accept=["text/plain", "application/pdf"],
@@ -101,36 +85,27 @@ async def on_chat_start():
 
     file = files[0]
 
-    msg = cl.Message(
-        content=f"Processing `{file.name}`..."
-    )
+    msg = cl.Message(content=f"Processing `{file.name}`...")
     await msg.send()
 
-    # load the file
     texts = process_file(file)
-
     print(f"Processing {len(texts)} text chunks")
 
-    # Create a dict vector store
     vector_db = VectorDatabase()
     vector_db = await vector_db.abuild_from_list(texts)
-    
-    # Initialize ChatOpenAI (API key will be automatically picked up from environment variable)
-    chat_openai = ChatOpenAI()
 
-    # Create a chain
+    chat_llm = HuggingFaceLLM()
+
     retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
         vector_db_retriever=vector_db,
-        llm=chat_openai
+        llm=chat_llm
     )
-    
-    # Let the user know that the system is ready
+
     msg.content = f"Processing `{file.name}` done. You can now ask questions!"
     await msg.update()
 
     cl.user_session.set("chain", retrieval_augmented_qa_pipeline)
 
-
 @cl.on_message
 async def main(message):
     chain = cl.user_session.get("chain")
@@ -141,4 +116,4 @@ async def main(message):
     async for stream_resp in result["response"]:
         await msg.stream_token(stream_resp)
 
-    await msg.send()
+    await msg.send()