Function get_response

taskbench/inference.py:216–293

(url, header, payload, id, reformat, reformat_by, dependency_type, log_detail=False)

    wf.flush()

async def get_response(url, header, payload, id, reformat, reformat_by, dependency_type, log_detail=False):
    async with aiohttp.ClientSession() as session:
        async with session.post(url, headers=header, data=payload, timeout=300) as response:
            resp = await response.json()

    if response.status == 429:
        raise RateLimitError(f"{resp}")
    if response.status != 200:
        raise Exception(f"{resp}")
    
    if log_detail:
        logger.info(json.loads(payload)["messages"][0]["content"])
        logger.info(resp["choices"][0]["message"]["content"])

    oring_content = resp["choices"][0]["message"]["content"]
    oring_content = oring_content.replace("\n", "")
    oring_content = oring_content.replace("\_", "_")
    content = oring_content.replace("\\", "")

    start_pos = content.find("RESULT #:")
    if start_pos!=-1:
        content = content[start_pos+len("RESULT #:"):]
        
    content = content[content.find("{"):content.rfind("}")+1]
    try:
        content = json.loads(content)
        if isinstance(content, list) and len(content):
            merge_content = {}
            for c in content:
                for k, v in c.items():
                    merge_content[k].extend(v) if k in merge_content else merge_content.update({k: v})
        return content
    except json.JSONDecodeError as e:
        if reformat:
            if dependency_type == "resource":
                prompt = """Please format the result # RESULT # to a strict JSON format # STRICT JSON FORMAT #. \nRequirements:\n1. Do not change the meaning of task steps and task nodes;\n2. Don't tolerate any possible irregular formatting to ensure that the generated content can be converted by json.loads();\n3. You must output the result in this schema: {"task_steps": [ step description of one or more steps ], "task_nodes": [{"task": "tool name must be from # TOOL LIST #", "arguments": [ a concise list of arguments for the tool. Either original text, or user-mentioned filename, or tag '<node-j>' (start from 0) to refer to the output of the j-th node. ]}]}\n# RESULT #:{{illegal_result}}\n# STRICT JSON FORMAT #:"""
            else:
                prompt = """Please format the result # RESULT # to a strict JSON format # STRICT JSON FORMAT #. \nRequirements:\n1. Do not change the meaning of task steps, task nodes and task links;\n2. Don't tolerate any possible irregular formatting to ensure that the generated content can be converted by json.loads();\n3. Pay attention to the matching of brackets. Write in a compact format and avoid using too many space formatting controls;\n4. You must output the result in this schema: {"task_steps": [ "concrete steps, format as Step x: Call xxx tool with xxx: 'xxx' and xxx: 'xxx'" ], "task_nodes": [{"task": "task name must be from # TASK LIST #", "arguments": [ {"name": "parameter name", "value": "parameter value, either user-specified text or the specific name of the tool whose result is required by this node"} ]}], "task_links": [{"source": "task name i", "target": "task name j"}]}\n# RESULT #:{{illegal_result}}\n# STRICT JSON FORMAT #:"""
            prompt = prompt.replace("{{illegal_result}}", oring_content)
            payload = json.loads(payload)
            if reformat_by != "self":
                payload["model"] = reformat_by

            if log_detail:
                logger.info(f"{emoji.emojize(':warning:')}  #id {id} Illegal JSON format: {content}")
                logger.info(f"{emoji.emojize(':sparkles:')} #id {id} Detected illegal JSON format, try to reformat by {payload['model']}...")

            payload["messages"][0]["content"] = prompt
            payload = json.dumps(payload)
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, headers=header, data=payload, timeout=120) as response:
                    resp = await response.json()

            if response.status == 429:
                raise RateLimitError(f"{resp}")
            if response.status != 200:
                raise Exception(f"{resp}")