Update query_solver.py

query_solver.py

@@ -3,8 +3,35 @@ import json
 import autogen
 from autogen.agentchat.contrib.math_user_proxy_agent import MathUserProxyAgent
 
+# Function to run the query
 def run_query(math_problem, api_key):
-    # ... (Same code as before)
+    config_list = [
+        {
+            'model': 'gpt-3.5-turbo',
+            'api_key': api_key,
+        },
+    ]
+
+    autogen.ChatCompletion.start_logging()
+
+    assistant = autogen.AssistantAgent(
+        name="assistant",
+        system_message="You are a helpful assistant.",
+        llm_config={
+            "request_timeout": 600,
+            "seed": 42,
+            "config_list": config_list,
+        }
+    )
+
+    mathproxyagent = MathUserProxyAgent(
+        name="mathproxyagent",
+        human_input_mode="NEVER",
+        code_execution_config={"use_docker": False},
+    )
+
+    return mathproxyagent.initiate_chat(assistant, problem=math_problem)
+   
 
 if __name__ == "__main__":
     input_data = json.loads(sys.stdin.read())