Update src/insurance_assistants/ui.py

#1
by Shamik - opened
Files changed (1) hide show
  1. src/insurance_assistants/ui.py +8 -4
src/insurance_assistants/ui.py CHANGED
@@ -15,7 +15,7 @@ from smolagents.models import ChatMessageStreamDelta
15
  from src.insurance_assistants.agents import manager_agent
16
  from src.insurance_assistants.consts import PRIMARY_HEADING, PROJECT_ROOT_DIR, PROMPT_PREFIX
17
 
18
- load_dotenv(override=True)
19
 
20
  # Setup logger
21
  logging.basicConfig(level=logging.INFO)
@@ -86,8 +86,10 @@ class UI:
86
  intermediate_text += step_log.content or ""
87
  yield intermediate_text
88
 
89
- def interact_with_agent(self, prompt, messages, session_state):
90
  # Get or create session-specific agent
 
 
91
  if "agent" not in session_state:
92
  # session_state["agent"] = CodeAgent(tools=[], model=InfenceClientModel())
93
  session_state["agent"] = manager_agent
@@ -220,6 +222,8 @@ class UI:
220
  """
221
  )
222
  with gr.Group():
 
 
223
  gr.Markdown(
224
  value="**Your question, please...**", container=True
225
  )
@@ -295,7 +299,7 @@ class UI:
295
  ).then(
296
  fn=self.interact_with_agent,
297
  # Include session_state in function calls
298
- inputs=[stored_messages, chatbot, session_state],
299
  outputs=[chatbot],
300
  ).then(
301
  fn=lambda: (
@@ -315,7 +319,7 @@ class UI:
315
  ).then(
316
  fn=self.interact_with_agent,
317
  # Include session_state in function calls
318
- inputs=[stored_messages, chatbot, session_state],
319
  outputs=[chatbot],
320
  ).then(
321
  fn=lambda: (
 
15
  from src.insurance_assistants.agents import manager_agent
16
  from src.insurance_assistants.consts import PRIMARY_HEADING, PROJECT_ROOT_DIR, PROMPT_PREFIX
17
 
18
+ # load_dotenv(override=True)
19
 
20
  # Setup logger
21
  logging.basicConfig(level=logging.INFO)
 
86
  intermediate_text += step_log.content or ""
87
  yield intermediate_text
88
 
89
+ def interact_with_agent(self, prompt, messages, session_state, api_key):
90
  # Get or create session-specific agent
91
+ if not api_key or not api_key.startswith("hf"):
92
+ raise ValueError("Incorrect API key")
93
  if "agent" not in session_state:
94
  # session_state["agent"] = CodeAgent(tools=[], model=InfenceClientModel())
95
  session_state["agent"] = manager_agent
 
222
  """
223
  )
224
  with gr.Group():
225
+ api_key = gr.Textbox(placeholder="Enter your HuggingFace Inference API KEY HERE",
226
+ label="API key",show_label=True, type='password')
227
  gr.Markdown(
228
  value="**Your question, please...**", container=True
229
  )
 
299
  ).then(
300
  fn=self.interact_with_agent,
301
  # Include session_state in function calls
302
+ inputs=[stored_messages, chatbot, session_state, api_key],
303
  outputs=[chatbot],
304
  ).then(
305
  fn=lambda: (
 
319
  ).then(
320
  fn=self.interact_with_agent,
321
  # Include session_state in function calls
322
+ inputs=[stored_messages, chatbot, session_state, api_key],
323
  outputs=[chatbot],
324
  ).then(
325
  fn=lambda: (