Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import openai, os, time, wandb
+import openai, os, time
 
 from dotenv import load_dotenv, find_dotenv
 from langchain.chains import LLMChain, RetrievalQA
@@ -14,8 +14,8 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.vectorstores import MongoDBAtlasVectorSearch
 from pymongo import MongoClient
+
 from trace import wandb_trace
-#from wandb.sdk.data_types.trace_tree import Trace
 
 _ = load_dotenv(find_dotenv())
 
@@ -36,8 +36,6 @@ MONGODB_INDEX_NAME        = "default"
 LLM_CHAIN_PROMPT = PromptTemplate(input_variables = ["question"], template = os.environ["LLM_TEMPLATE"])
 RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], template = os.environ["RAG_TEMPLATE"])
 
-#WANDB_API_KEY = os.environ["WANDB_API_KEY"]
-
 RAG_OFF     = "Off"
 RAG_CHROMA  = "Chroma"
 RAG_MONGODB = "MongoDB"
@@ -116,43 +114,6 @@ def rag_chain(llm, prompt, db):
     completion = rag_chain({"query": prompt})
     return completion, rag_chain
 
-#def wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms):
-#    wandb.init(project = "openai-llm-rag")
-#    
-#    trace = Trace(
-#        kind = "chain",
-#        name = "" if (chain == None) else type(chain).__name__,
-#        status_code = "success" if (str(err_msg) == "") else "error",
-#        status_message = str(err_msg),
-#        metadata = {"chunk_overlap": "" if (rag_option == RAG_OFF) else config["chunk_overlap"],
-#                    "chunk_size": "" if (rag_option == RAG_OFF) else config["chunk_size"],
-#                   } if (str(err_msg) == "") else {},
-#        inputs = {"rag_option": rag_option,
-#                  "prompt": prompt,
-#                  "chain_prompt": (str(chain.prompt) if (rag_option == RAG_OFF) else 
-#                                   str(chain.combine_documents_chain.llm_chain.prompt)),
-#                  "source_documents": "" if (rag_option == RAG_OFF) else str([doc.metadata["source"] for doc in completion["source_documents"]]),
-#                 } if (str(err_msg) == "") else {},
-#        outputs = {"result": result,
-#                   "generation_info": str(generation_info),
-#                   "llm_output": str(llm_output),
-#                   "completion": str(completion),
-#                  } if (str(err_msg) == "") else {},
-#        model_dict = {"client": (str(chain.llm.client) if (rag_option == RAG_OFF) else
-#                                 str(chain.combine_documents_chain.llm_chain.llm.client)),
-#                      "model_name": (str(chain.llm.model_name) if (rag_option == RAG_OFF) else
-#                                     str(chain.combine_documents_chain.llm_chain.llm.model_name)),
-#                      "temperature": (str(chain.llm.temperature) if (rag_option == RAG_OFF) else
-#                                      str(chain.combine_documents_chain.llm_chain.llm.temperature)),
-#                      "retriever": ("" if (rag_option == RAG_OFF) else str(chain.retriever)),
-#                     } if (str(err_msg) == "") else {},
-#        start_time_ms = start_time_ms,
-#        end_time_ms = end_time_ms
-#    )
-#    
-#    trace.log("evaluation")
-#    wandb.finish()
-
 def invoke(openai_api_key, rag_option, prompt):
     if (openai_api_key == ""):
         raise gr.Error("OpenAI API Key is required.")
@@ -199,14 +160,25 @@ def invoke(openai_api_key, rag_option, prompt):
             llm_output = completion.llm_output
     except Exception as e:
         err_msg = e
+
         raise gr.Error(e)
     finally:
         end_time_ms = round(time.time() * 1000)
         
-        wandb_trace(rag_option, prompt, completion, result, generation_info, llm_output, chain, err_msg, start_time_ms, end_time_ms)
+        wandb_trace(rag_option, 
+                    prompt, 
+                    completion, 
+                    result, 
+                    generation_info, 
+                    llm_output, 
+                    chain, 
+                    err_msg, 
+                    start_time_ms, 
+                    end_time_ms)
     return result
 
 gr.close_all()
+
 demo = gr.Interface(fn=invoke, 
                     inputs = [gr.Textbox(label = "OpenAI API Key", type = "password", lines = 1), 
                               gr.Radio([RAG_OFF, RAG_CHROMA, RAG_MONGODB], label = "Retrieval Augmented Generation", value = RAG_OFF),
@@ -215,4 +187,5 @@ demo = gr.Interface(fn=invoke,
                     outputs = [gr.Textbox(label = "Completion", lines = 1)],
                     title = "Generative AI - LLM & RAG",
                     description = os.environ["DESCRIPTION"])
+
 demo.launch()