# qwen_client_v2.py
import requests
import json

class QwenClient:
    def __init__(self, endpoint_url):
        self.endpoint_url = endpoint_url

    def generate_response(self, prompt):
        try:
            headers = {
                'Content-Type': 'application/json',
            }
            data = {
                'prompt': prompt
            }
            response = requests.post(self.endpoint_url, headers=headers, data=json.dumps(data), stream=False)

            if response.status_code == 200:
                return response.json()['response']
            else:
                return f"Error: {response.status_code} - {response.text}"
        except Exception as e:
            return f"Exception: {e}"


if __name__ == '__main__':
    # Example Usage (Replace with your actual endpoint)
    endpoint = 'http://localhost:11434/api/generate'
    client = QwenClient(endpoint)

    test_prompt = "Write a short poem about the beauty of artificial intelligence."
    response = client.generate_response(test_prompt)

    print(f"Prompt: {test_prompt}\nResponse: {response}")

    test_prompt2 = "Summarize the plot of Hamlet in 3 sentences."
    response2 = client.generate_response(test_prompt2)

    print(f"Prompt: {test_prompt2}\nResponse: {response2}")