import logging
import os

import openai
from langchain.chat_models import AzureChatOpenAI

import vishwa
from vishwa.mlmonitor.langchain.decorators.map_xpuls_project import MapXpulsProject
from vishwa.mlmonitor.langchain.decorators.telemetry_override_labels import TelemetryOverrideLabels
from vishwa.mlmonitor.langchain.instrument import LangchainTelemetry
from vishwa.mlmonitor.langchain.patches.xp_prompt_template import XPChatPromptTemplate
from vishwa.prompt_hub import PromptClient

logger = logging.getLogger(__name__)

openai.api_key = os.getenv("OPENAI_API_KEY")
openai.api_type = "azure"
openai.api_base = os.getenv("OPENAI_URL")
os.environ["OPENAI_API_BASE"] = os.getenv("OPENAI_URL")
os.environ["OPENAI_API_VERSION"] = "2023-03-15-preview"
openai.api_version = "2023-03-15-preview"

# Set this to enable Advanced prompt tracing with server


default_labels = {"system": "openai-ln-test", "agent_name": "fallback_value"}
vishwa.host_url = "https://test-api.vishwa.ai"
vishwa.api_key = "****************************************"
vishwa.adv_tracing_enabled = "true"

LangchainTelemetry(
    default_labels=default_labels,
).auto_instrument()

chat_model = AzureChatOpenAI(
    deployment_name="gpt35turbo",
    model_name="gpt-35-turbo",
    temperature=0
)

prompt_client = PromptClient(
    prompt_id="clrfm4v70jnlb1kph240",
    environment_name="dev"
)
@TelemetryOverrideLabels(agent_name="chat_agent_alpha")
@MapXpulsProject(project_id="defaultoPIt9USSR")  # Get Project ID from console
def run_openai_agent():
    # prompt = ChatPromptTemplate.from_template("tell me a joke about {foo}")
    data = prompt_client.get_prompt({"variable-1": "I'm the first variable"})
    prompt = XPChatPromptTemplate.from_template(data)
    chain = prompt | chat_model
    try:
        res = chain.invoke({"foo": "bears"})
    except ValueError as e:
        res = str(e)
        if not res.startswith("Could not parse LLM output: `"):
            raise e
        logger.error(f" Got ValueError: {e}")
        res = res.removeprefix("Could not parse LLM output: `").removesuffix("`")

    return res

langchain.chat_models.AzureChatOpenAI
[((498, 525), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (515, 525), False, 'import logging\n'), ((544, 571), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (553, 571), False, 'import os\n'), ((616, 639), 'os.getenv', 'os.getenv', (['"""OPENAI_URL"""'], {}), "('OPENAI_URL')\n", (625, 639), False, 'import os\n'), ((672, 695), 'os.getenv', 'os.getenv', (['"""OPENAI_URL"""'], {}), "('OPENAI_URL')\n", (681, 695), False, 'import os\n'), ((1165, 1256), 'langchain.chat_models.AzureChatOpenAI', 'AzureChatOpenAI', ([], {'deployment_name': '"""gpt35turbo"""', 'model_name': '"""gpt-35-turbo"""', 'temperature': '(0)'}), "(deployment_name='gpt35turbo', model_name='gpt-35-turbo',\n    temperature=0)\n", (1180, 1256), False, 'from langchain.chat_models import AzureChatOpenAI\n'), ((1284, 1354), 'vishwa.prompt_hub.PromptClient', 'PromptClient', ([], {'prompt_id': '"""clrfm4v70jnlb1kph240"""', 'environment_name': '"""dev"""'}), "(prompt_id='clrfm4v70jnlb1kph240', environment_name='dev')\n", (1296, 1354), False, 'from vishwa.prompt_hub import PromptClient\n'), ((1366, 1420), 'vishwa.mlmonitor.langchain.decorators.telemetry_override_labels.TelemetryOverrideLabels', 'TelemetryOverrideLabels', ([], {'agent_name': '"""chat_agent_alpha"""'}), "(agent_name='chat_agent_alpha')\n", (1389, 1420), False, 'from vishwa.mlmonitor.langchain.decorators.telemetry_override_labels import TelemetryOverrideLabels\n'), ((1422, 1468), 'vishwa.mlmonitor.langchain.decorators.map_xpuls_project.MapXpulsProject', 'MapXpulsProject', ([], {'project_id': '"""defaultoPIt9USSR"""'}), "(project_id='defaultoPIt9USSR')\n", (1437, 1468), False, 'from vishwa.mlmonitor.langchain.decorators.map_xpuls_project import MapXpulsProject\n'), ((1693, 1733), 'vishwa.mlmonitor.langchain.patches.xp_prompt_template.XPChatPromptTemplate.from_template', 'XPChatPromptTemplate.from_template', (['data'], {}), '(data)\n', (1727, 1733), False, 'from vishwa.mlmonitor.langchain.patches.xp_prompt_template import XPChatPromptTemplate\n'), ((1076, 1125), 'vishwa.mlmonitor.langchain.instrument.LangchainTelemetry', 'LangchainTelemetry', ([], {'default_labels': 'default_labels'}), '(default_labels=default_labels)\n', (1094, 1125), False, 'from vishwa.mlmonitor.langchain.instrument import LangchainTelemetry\n')]