import os
import litellm

class GenesisClient:
    def __init__(self, model_config=None):
        self.model_config = model_config or {}
        self.default_model = os.environ.get("LITELLM_DEFAULT_MODEL", "gpt-3.5-turbo")  # Default model if none is specified

    def chat_completion(self, messages, model=None, **kwargs):
        model = model or self.default_model
        
        # Apply model-specific configuration if available
        config = self.model_config.get(model, {})
        kwargs.update(config)

        try:
            response = litellm.completion(
                model=model,
                messages=messages,
                **kwargs
            )
            return response
        except Exception as e:
            print(f"Error during chat completion: {e}")
            raise

    def embedding(self, input, model=None, **kwargs):
        model = model or self.default_model

        # Apply model-specific configuration if available
        config = self.model_config.get(model, {})
        kwargs.update(config)

        try:
            response = litellm.embedding(
                model=model,
                input=input,
                **kwargs
            )
            return response
        except Exception as e:
            print(f"Error during embedding: {e}")
            raise


if __name__ == '__main__':
    # Example Usage
    client = GenesisClient(model_config={
        "gpt-3.5-turbo": {
            "temperature": 0.7
        },
        "claude-v1.3": {
            "max_tokens": 500
        }
    })

    messages = [{"role": "user", "content": "Tell me a joke."}]

    try:
        response = client.chat_completion(messages=messages, model="gpt-3.5-turbo")
        print(f"GPT-3.5-turbo Response: {response}")

        response = client.chat_completion(messages=messages, model="claude-v1.3")
        print(f"Claude Response: {response}")

        embedding_response = client.embedding(input="This is a test embedding.", model="text-embedding-ada-002")
        print(f"Embedding Response: {embedding_response}")

    except Exception as e:
        print(f"An error occurred: {e}")