import os
from typing import Optional
import instructor
from openai import OpenAI
from pydantic import BaseModel, Field

class InstructModel(BaseModel):
    completion: str = Field(..., title="Completion", description="The completion generated by the model.")

class Instruct:
    def __init__(self, api_key: Optional[str] = None):
        """
        Constructs all the necessary attributes for the Instruct object.

        Parameters
        ----------
        api_key : str, optional
            The OpenAI API key to use (default is None)

        Raises
        ------
        ValueError
            If the API key is not available
        """
        self.model = None
        self.query = None
        self.system = None

        if api_key is None:
            api_key = os.getenv("OPENAI_API_KEY")

        if api_key is None:
            raise ValueError("OpenAI API key is not available")

        self.llm = instructor.patch(OpenAI(api_key=api_key))

    def instruct(
        self,
        system: Optional[str] = None,
        query: Optional[str] = None,
        model: str = "gpt-3.5-turbo-1106",
        pydantic_model: BaseModel = InstructModel,
    ):
        """
        Lets the user create their own Pydantic model and use it, if they choose.
        """
        if query is None:
            raise ValueError("Query is required for instruct()")

        if model == "3":
            model = "gpt-3.5-turbo-1106"
        elif model == "4":
            model = "gpt-4-turbo-preview"

        completion = self.llm.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": query},
            ],
            response_model=pydantic_model,
        )

        return completion

if __name__ == "__main__":
    instruct = Instruct()
    
    system_prompt = "You are a helpful assistant."
    user_query = "What is the capital of France?"

    completion = instruct.instruct(system=system_prompt, query=user_query)
    print(completion.completion)