2323target_source_chunks = int (os .environ .get ("TARGET_SOURCE_CHUNKS" , 4 ))
2424loaded_model_lock = threading .Lock ()
2525loaded_model = threading .Event ()
26- logger = logging .getLogger (__name__ )
27- logger .setLevel (logging .DEBUG )
26+ app .logger .setLevel (logging .DEBUG )
2827
2928
3029def document_loader ():
3130 try :
3231 load_dir = "retrieval"
33- logger .debug ("Loading documents from %s" , load_dir )
32+ app . logger .debug ("Loading documents from %s" , load_dir )
3433 loader = DirectoryLoader (
3534 load_dir ,
3635 exclude = ["**/*.png" , "**/images/**" , "**/images/*" , "**/*.pdf" ],
3736 recursive = True ,
3837 loader_cls = UnstructuredMarkdownLoader ,
3938 )
4039 documents = loader .load ()
41- logger .debug ("Loaded %s documents" , len (documents ))
40+ app . logger .debug ("Loaded %s documents" , len (documents ))
4241 text_splitter = CharacterTextSplitter (chunk_size = 1000 , chunk_overlap = 100 )
4342 texts = text_splitter .split_documents (documents )
4443 embeddings = get_embeddings ()
@@ -48,7 +47,7 @@ def document_loader():
4847 retriever = db .as_retriever (search_kwargs = {"k" : target_source_chunks })
4948 return retriever
5049 except Exception as e :
51- logger .error ("Error loading documents %s" , e , exc_info = True )
50+ app . logger .error ("Error loading documents %s" , e , exc_info = True )
5251 raise e
5352
5453
@@ -93,32 +92,37 @@ def qa_app(qa, query):
9392
9493@app .route ("/chatbot/genai/init" , methods = ["POST" ])
9594def init_bot ():
95+ app .logger .debug ("Initializing bot" )
9696 try :
9797 with loaded_model_lock :
98- if "openai_api_key" in request .json :
99- logger .debug ("Initializing bot %s" , request .json ["openai_api_key" ])
100- os .environ ["OPENAI_API_KEY" ] = request .json ["openai_api_key" ]
101- global vulnerable_app_qa , retriever
102- retriever = document_loader ()
103- llm = get_llm ()
104- vulnerable_app_qa = get_qa_chain (llm , retriever )
105- loaded_model .set ()
106- return jsonify ({"message" : "Model Initialized" }), 200
107- else :
98+ if os .environ .get ("CHATBOT_OPENAI_API_KEY" ) is not None :
99+ app .logger .info ("Using OpenAI API Key from environment" )
100+ os .environ ["OPENAI_API_KEY" ] = os .environ .get ("CHATBOT_OPENAI_API_KEY" )
101+ elif "openai_api_key" not in request .json :
108102 return jsonify ({"message" : "openai_api_key not provided" }, 400 )
103+ app .logger .debug ("Initializing bot %s" , request .json ["openai_api_key" ])
104+ os .environ ["OPENAI_API_KEY" ] = request .json ["openai_api_key" ]
105+ global vulnerable_app_qa , retriever
106+ retriever = document_loader ()
107+ llm = get_llm ()
108+ vulnerable_app_qa = get_qa_chain (llm , retriever )
109+ loaded_model .set ()
110+ return jsonify ({"message" : "Model Initialized" }), 200
111+
109112 except Exception as e :
110- logger .error ("Error initializing bot " , e )
111- logger .debug ("Error initializing bot " , e , exc_info = True )
113+ app . logger .error ("Error initializing bot " , e )
114+ app . logger .debug ("Error initializing bot " , e , exc_info = True )
112115 return jsonify ({"message" : "Not able to initialize model " + str (e )}), 400
113116
114117
115118@app .route ("/chatbot/genai/state" , methods = ["GET" ])
116119def state_bot ():
120+ app .logger .debug ("Checking state" )
117121 try :
118122 if loaded_model .is_set ():
119123 return jsonify ({"initialized" : "true" , "message" : "Model already loaded" })
120124 except Exception as e :
121- logger .error ("Error checking state " , e )
125+ app . logger .error ("Error checking state " , e )
122126 return jsonify ({"message" : "Error checking state " + str (e )}), 200
123127 return (
124128 jsonify ({"initialized" : "false" , "message" : "Model needs to be initialized" }),
@@ -128,15 +132,20 @@ def state_bot():
128132
129133@app .route ("/chatbot/genai/ask" , methods = ["POST" ])
130134def ask_bot ():
135+ app .logger .debug ("Asking bot" )
131136 question = request .json ["question" ]
132137 global vulnerable_app_qa
133138 answer = qa_app (vulnerable_app_qa , question )
134- logger .info ("###########################################" )
135- logger .info ("Test Attacker Question: %s" , question )
136- logger .info ("Vulnerability App Answer: %s" , answer )
137- logger .info ("###########################################" )
139+ app . logger .info ("###########################################" )
140+ app . logger .info ("Test Attacker Question: %s" , question )
141+ app . logger .info ("Vulnerability App Answer: %s" , answer )
142+ app . logger .info ("###########################################" )
138143 return jsonify ({"answer" : answer }), 200
139144
140145
141146if __name__ == "__main__" :
142147 app .run (host = "0.0.0.0" , port = 5002 , debug = True )
148+ else :
149+ gunicorn_logger = logging .getLogger ("gunicorn.error" )
150+ app .logger .handlers = gunicorn_logger .handlers
151+ app .logger .setLevel (logging .DEBUG )
0 commit comments