Refactor Agent initialization and enhance respond function to support web search; add check_format utility function.

app.py

@@ -47,6 +47,15 @@ def load_file(path: str) -> list | dict:
     else:
         return {"raw document text": text, "file path": path}
     
+def check_format(answer: str | list, *args, **kwargs ) -> list:
+    """ Check if the answer is a list. """
+    
+    print("Checking format of the answer:", answer)
+    if isinstance(answer, list):
+        return answer
+    else:
+        return [answer]
+    
 
 ## tools definition
 @tool
@@ -134,13 +143,15 @@ def generate_image(prompt: str, neg_prompt: str) -> Image.Image:
 ## agent definition
 class Agent:
     def __init__(self, ):
-        client = HfApiModel("Qwen/Qwen3-32B", provider="nebius", api_key=os.getenv("NEBIUS_API_KEY"))
+        client = HfApiModel("google/gemma-3-27b-it", provider="nebius", api_key=os.getenv("NEBIUS_API_KEY"))
         self.agent = CodeAgent(
             model=client,
             tools=[DuckDuckGoSearchTool(max_results=5), VisitWebpageTool(max_output_length=20000), generate_image, download_images, transcribe_audio],
             additional_authorized_imports=["pandas", "PIL", "io"],
             planning_interval=1,
-            max_steps=8,
+            max_steps=5,
+            stream_outputs=False,
+            final_answer_checks=[check_format]
         )
         with open("system_prompt.txt", "r") as f:
             system_prompt = f.read()
@@ -156,14 +167,17 @@ class Agent:
         return answer
 
 ## gradio functions
-def respond(message: str, history : dict):
+def respond(message: str, history : dict, web_search: bool = False):
     
     # input
     print("history:", history)
     text = message.get("text", "")
-    if not message.get("files"): # no files uploaded
+    if not message.get("files") and not web_search: # no files uploaded
         print("No files received.")
-        message = agent(text, conversation_history=history) # conversation_history is a dict with the history of the conversation
+        message = agent(text + "\nADDITIONAL CONTRAINT: Don't use web search", conversation_history=history) # conversation_history is a dict with the history of the conversation 
+    elif not message.get("files") and web_search==True: # no files uploaded
+        print("No files received + web search enabled.")
+        message = agent(text, conversation_history=history)
     else:
         files = message.get("files", [])
         print(f"files received: {files}")
@@ -173,12 +187,10 @@ def respond(message: str, history : dict):
         else:
             file = load_file(files[0])
             message = agent(text, files=file, conversation_history=history)
-
     # output
     print("Agent response:", message)
     
     return message
-    
 
 def initialize_agent():
     agent = Agent()
@@ -196,9 +208,13 @@ with gr.Blocks() as demo:
                         title='MultiAgent System for Screenplay Creation and Editing',
                         show_progress='full',
                         fill_height=True,
-                        fill_width=True,
+                        fill_width=False,
                         save_history=True,
-                    )
+                        additional_inputs=[
+                        gr.Checkbox(value=False, label="Web Search", 
+                                    info="Enable web search to find information online. If disabled, the agent will only use the provided files and images."
+                                    ),
+                    ])
 
 
 if __name__ == "__main__":