Update ai_assistant.py

ai_assistant.py

@@ -1,3 +1,4 @@
+import os
 from langchain.document_loaders.csv_loader import CSVLoader
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.embeddings import CacheBackedEmbeddings
@@ -5,49 +6,51 @@ from langchain_community.vectorstores import FAISS
 from langchain.storage import LocalFileStore
 from langchain.chains import RetrievalQA
 from langchain_openai import ChatOpenAI
-import os
 
 def create_index():
-    # load the data
-    dir = os.path.dirname(__file__)    
-    df_path = dir + '/data/train.csv'
-    loader = CSVLoader(file_path = df_path)
-    data = loader.load()
+    # Load the data from CSV file
+    data_loader = CSVLoader(file_path="train.csv")
+    data = data_loader.load()
     
-    # create the embeddings model
+    # Create the embeddings model
     embeddings_model = OpenAIEmbeddings()
 
-    # create the cache backed embeddings in vector store
+    # Create the cache backed embeddings in vector store
     store = LocalFileStore("./cache")
-    cached_embeder = CacheBackedEmbeddings.from_bytes_store(
+    cached_embedder = CacheBackedEmbeddings.from_bytes_store(
         embeddings_model, store, namespace=embeddings_model.model
     )
+    
+    # Create FAISS vector store from documents
     vector_store = FAISS.from_documents(data, embeddings_model)
 
     return vector_store.as_retriever()
 
-def setup(openai_key):
-     # Set the API key for OpenAI
+def setup_openai(openai_key):
+    # Set the API key for OpenAI
     os.environ["OPENAI_API_KEY"] = openai_key
-    retriver = create_index()
-    llm = ChatOpenAI(model="gpt-4")
-    return retriver, llm
-
-def ai_doctor(openai_key,query):
+    
+    # Create index retriever
+    retriever = create_index()
+    
+    # Initialize ChatOpenAI model
+    chat_openai_model = ChatOpenAI(model="gpt-4")
+    
+    return retriever, chat_openai_model
 
-    # Setup
-    retriever,llm = setup(openai_key)
+def ai_doctor_chat(openai_key, query):
+    # Setup OpenAI environment
+    retriever, chat_model = setup_openai(openai_key)
+    
     # Create the QA chain
     handler = StdOutCallbackHandler()
-
     qa_with_sources_chain = RetrievalQA.from_chain_type(
-        llm=llm,
+        llm=chat_model,
         retriever=retriever,
         callbacks=[handler],
         return_source_documents=True
     )
 
-    # Ask a question
-    res = qa_with_sources_chain({"query":query})
-    return (res['result'])
-
+    # Ask a question/query
+    res = qa_with_sources_chain({"query": query})
+    return res['result']