import os
from athina_client.prompt import Prompt, Slug
from athina_client.keys import AthinaApiKey

AthinaApiKey.set_key(os.getenv('ATHINA_API_KEY'))

Prompt.create_prompt(
    slug='test-staging',
    prompt=[{
        "role": "system",
        "content": "You are an AI that answers questions in less than 50 words"
    },
    {
        "role": "user",
        "content": "what does {{company}} does?"
    }],
    model="gpt-4o",
    commit_message="Initial prompt commit",
    parameters={
        "temperature": 0.5
    }
)


Prompt.run_prompt(
    slug='test-staging',
    # the following fields are optional
    version=2,
    model="gpt-4o",
    variables={
        "company": "NVIDIA"
    },
    parameters={
        "temperature": 1,
        "max_tokens": 1000
    },
)

from athina_logger.api_key import AthinaApiKey

AthinaApiKey.set_api_key(os.getenv('ATHINA_API_KEY'))

response = client.chat.completions.create(
      model='gpt-4-1106-preview',
      messages=[{"role": "user", "content": "What is machine learning?"}],
  )

response = response.model_dump() # For openai > 1 version

try:
  InferenceLogger.log_inference(
      prompt_slug="sdk_test",
      prompt=messages,
      language_model_id="gpt-4-1106-preview",
      response=response,
      external_reference_id="abc",
      cost=0.0123,
      custom_attributes={
          "name": "John Doe"
          # Your custom attributes
      }
  )
except Exception as e:
  if isinstance(e, CustomException):
      print(e.status_code)
      print(e.message)
  else:
      print(e)

import os
from athina_client.keys import AthinaApiKey

AthinaApiKey.set_key(os.getenv('ATHINA_API_KEY'))

from athina_client.datasets import Dataset

try:
  dataset = Dataset.create(
    name='test_dataset',
    description="Optional description", # optional
    language_model_id="gpt-4", # optional,
    rows=[
      { "query": "Who let the dogs out?", "response": "Who, who, who, who, who?" },
      ...
    ]
  )
except Exception as e:
  print(f"Failed to create dataset: {e}")

query GetPromptRuns($limit: Int!, $page: Int!) {
  getPromptRunsByFilters(limit: $limit, page: $page) {
    id
    org_id
    prompt_slug
    language_model_id
    prompt_response
    prompt_tokens
  }
}