from openai import OpenAI
from typing import Optional, Dict, Any

class AIAssistant:
    """
    A wrapper class for consistent LLM API interactions.

    This class provides:
    - Unified interface for different LLM providers
    - Consistent handling of generation parameters
    - Support for streaming responses

    Attributes:
        client: Initialized API client (OpenAI, Anthropic, etc.)
        model: Name of the model to use
    """
    def __init__(self, client: OpenAI, model: str):
        self.client = client
        self.model = model

    def generate_response(self,
                         prompt_template: Any,
                         generation_params: Optional[Dict] = None,
                         stream: bool = False,
                         **kwargs):
        """
        Generate LLM response using pthe rovided template and parameters.

        Args:
            prompt_template: Template object with format method
            generation_params: Optional generation parameters
            stream: Whether to stream the response
            **kwargs: Variables for prompt template

        Returns:
            API response object or streamed response

        Example:
            assistant.generate_response(
                prompt_template=template,
                temperature=0.7,
                topic="AI safety"
            )
        """
        messages = prompt_template.format(**kwargs)
        params = generation_params or {}

        completion = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            stream=stream,
            **params
        )

        if stream:
            for chunk in completion:
                if chunk.choices[0].delta.content is not None:
                    print(chunk.choices[0].delta.content, end="")
            return completion

        return completion