(system_prompt: str, initial_query: str, client, model: str)
| 65 | return any(keyword in query.lower() for keyword in keywords) |
| 66 | |
| 67 | def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]: |
| 68 | query, request_code = extract_python_code(initial_query)[0] if extract_python_code(initial_query) else (initial_query, "") |
| 69 | |
| 70 | if should_execute_request_code(query) and request_code: |
| 71 | # Execute code from the request |
| 72 | code_output = execute_code(request_code) |
| 73 | context = f"Query: {query}\nCode:\n```python\n{request_code}\n```\nOutput:\n{code_output}" |
| 74 | |
| 75 | messages = [ |
| 76 | {"role": "system", "content": system_prompt}, |
| 77 | {"role": "user", "content": context} |
| 78 | ] |
| 79 | |
| 80 | response = client.chat.completions.create( |
| 81 | model=model, |
| 82 | messages=messages, |
| 83 | ) |
| 84 | |
| 85 | return response.choices[0].message.content.strip(), response.usage.completion_tokens |
| 86 | else: |
| 87 | # Get initial response from the model |
| 88 | messages = [ |
| 89 | {"role": "system", "content": system_prompt + EXECUTE_CODE_PROMPT} , |
| 90 | {"role": "user", "content": initial_query} |
| 91 | ] |
| 92 | |
| 93 | response = client.chat.completions.create( |
| 94 | model=model, |
| 95 | messages=messages, |
| 96 | ) |
| 97 | |
| 98 | initial_response = response.choices[0].message.content.strip() |
| 99 | response_code = extract_python_code(initial_response) |
| 100 | |
| 101 | if response_code: |
| 102 | # Execute code from the response |
| 103 | code_output = execute_code(response_code[0]) |
| 104 | context = f"Initial response:\n{initial_response}\n\nCode output:\n{code_output}" |
| 105 | |
| 106 | messages.append({"role": "assistant", "content": initial_response}) |
| 107 | messages.append({"role": "user", "content": f"Based on the code execution output, please provide a final response:\n{context}"}) |
| 108 | |
| 109 | final_response = client.chat.completions.create( |
| 110 | model=model, |
| 111 | messages=messages, |
| 112 | ) |
| 113 | |
| 114 | return final_response.choices[0].message.content.strip(), response.usage.completion_tokens + final_response.usage.completion_tokens |
| 115 | else: |
| 116 | return initial_response, response.usage.completion_tokens |
nothing calls this directly
no test coverage detected