from typing import Iterator, Optional
import openai
import json

from .llm_engine import LlmEngine

class OpenAILlmEngine(LlmEngine):
    """
    LLM Engine implementation using OpenAI's API.
    Supports streaming responses from chat completion models.
    """
    
    def __init__(
        self,
        model: str,
        api_key: str,
    ):
        """
        Initialize the OpenAI LLM Engine.

        Args:
            model: OpenAI model to use (default: gpt-4)
            api_key: OpenAI API key. If None, will try to read from OPENAI_API_KEY env var
        """
        self._model = model
        
        # Initialize OpenAI client
        self.client = openai.Client(
            api_key=api_key,
        )
                        
    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
        """
        Run inference using the system prompt and main context.

        Args:
            system_prompt: The system prompt string
            main_context: The main context string after templating

        Returns:
            Iterator[str]: An iterator that yields the generated text in chunks.
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context}
        ]
        
        stream = self.client.chat.completions.create(
            model=self._model,
            messages=messages,
            stream=True,
            temperature=0.3,
        )

        for chunk in stream:
            if content := chunk.choices[0].delta.content:
                yield content