query_solver.py

import sys
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):
    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())
    math_problem = input_data['math_problem']
    api_key = input_data['api_key']
    result = run_query(math_problem, api_key)
    print(result)