Switch to openai

app.py

@@ -11,7 +11,8 @@ model_config = ConfigDict(protected_namespaces=())
 setting_keys = gr.Interface(
     fn=set_keys,
     inputs=[
-        gr.Textbox(label="Enter your Cohere API key"),
+        #gr.Textbox(label="Enter your CO_API_KEY"),
+        gr.Textbox(label="Enter your OPENAI_API_KEY"),
         gr.Textbox(label="Enter your LLAMA_CLOUD_API_KEY"),
     ],
     outputs=gr.Textbox(label="Status")

main.py

@@ -2,12 +2,12 @@ from utils.retriever import get_query_engine
 from utils.index import create_index
 from utils.constant import INDEX_PATH, DATA_PATH
 import os
-from pathlib import Path
 import shutil
 
-def set_keys(co_api_key: str, llama_cloud_api_key: str) -> str:
+def set_keys(openai_api_key: str, llama_cloud_api_key: str) -> str:
     try:
-        os.environ["CO_API_KEY"] = co_api_key
+        #os.environ["CO_API_KEY"] = co_api_key
+        os.environ["OPENAI_API_KEY"] = openai_api_key
         os.environ["LLAMA_CLOUD_API_KEY"] = llama_cloud_api_key
         return "Keys are set successfully"
     

utils/constant.py

@@ -3,5 +3,3 @@ INDEX_PATH = "./index"
 
 TOP_K_RETRIEVAL = 10  
 TOP_N_RERANKER = 3
-
-EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"

utils/retriever.py

@@ -30,7 +30,7 @@ class QueryEngineManager:
         configure_settings()
         self.index = load_index(path=INDEX_PATH)
         self.nodes = list(self.index.docstore.docs.values())
-        self.reranker = LLMRerank(top_n=TOP_K_RETRIEVAL, choice_batch_size=1)
+        self.reranker = LLMRerank(top_n=TOP_N_RERANKER)
         
     def get_engine(self, bm25: bool = False, semantic: bool = False):
         if bm25:

utils/settings.py

@@ -1,5 +1,6 @@
 from llama_index.core import Settings
-from llama_index.llms.cohere import Cohere
+#from llama_index.llms.cohere import Cohere
+from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.core.node_parser import SemanticSplitterNodeParser
 
@@ -15,9 +16,14 @@ def load_llm():
     - Do not include information from external sources not provided by the user.
     '''
     
-    llm = Cohere(
-        system_prompt=system_prompt,
+    # llm = Cohere(
+    #     system_prompt=system_prompt,
         
+    # )
+    
+    llm = OpenAI(
+        model = "gpt-4o-mini",
+        system_prompt=system_prompt
     )
     
     return llm