import getpass
import os

# Make sure environment variables are set (or pass them directly to ChatRunPod)
if"RUNPOD_API_KEY"notin os.environ:
    os.environ["RUNPOD_API_KEY"]= getpass.getpass("Enter your RunPod API Key: ")

if"RUNPOD_ENDPOINT_ID"notin os.environ:
    os.environ["RUNPOD_ENDPOINT_ID"]=input(
"Enter your RunPod Endpoint ID (used if RUNPOD_CHAT_ENDPOINT_ID is not set): "
)

# Optionally use a different endpoint ID specifically for chat models
# if "RUNPOD_CHAT_ENDPOINT_ID" not in os.environ:
#     os.environ["RUNPOD_CHAT_ENDPOINT_ID"] = input("Enter your RunPod Chat Endpoint ID (Optional): ")

chat_endpoint_id= os.environ.get(
"RUNPOD_CHAT_ENDPOINT_ID", os.environ.get("RUNPOD_ENDPOINT_ID")
)
ifnot chat_endpoint_id:
raise ValueError(
"No RunPod Endpoint ID found. Please set RUNPOD_ENDPOINT_ID or RUNPOD_CHAT_ENDPOINT_ID."
)

from langchain_runpodimport ChatRunPod

chat= ChatRunPod(
    runpod_endpoint_id=chat_endpoint_id,# Specify the correct endpoint ID
    model_kwargs={
"max_new_tokens":512,
"temperature":0.7,
"top_p":0.9,
# Add other parameters supported by your endpoint handler
},
# Optional: Adjust polling
# poll_interval=0.2,
# max_polling_attempts=150
)

from langchain_core.messagesimport HumanMessage, SystemMessage

messages=[
    SystemMessage(content="You are a helpful AI assistant."),
    HumanMessage(content="What is the RunPod Serverless API flow?"),
]

# Invoke (Sync)
try:
    response= chat.invoke(messages)
print("--- Sync Invoke Response ---")
print(response.content)
except Exceptionas e:
print(
f"Error invoking Chat Model:{e}. Ensure endpoint ID/API key are correct and endpoint is active/compatible."
)

# Stream (Sync, simulated via polling /stream)
print("\n--- Sync Stream Response ---")
try:
for chunkin chat.stream(messages):
print(chunk.content, end="", flush=True)
print()# Newline
except Exceptionas e:
print(
f"\nError streaming Chat Model:{e}. Ensure endpoint handler supports streaming output format."
)

### Async Usage

# AInvoke (Async)
try:
    async_response=await chat.ainvoke(messages)
print("--- Async Invoke Response ---")
print(async_response.content)
except Exceptionas e:
print(f"Error invoking Chat Model asynchronously:{e}.")

# AStream (Async)
print("\n--- Async Stream Response ---")
try:
asyncfor chunkin chat.astream(messages):
print(chunk.content, end="", flush=True)
print()# Newline
except Exceptionas e:
print(
f"\nError streaming Chat Model asynchronously:{e}. Ensure endpoint handler supports streaming output format.\n"
)

from langchain_core.output_parsersimport StrOutputParser
from langchain_core.promptsimport ChatPromptTemplate

prompt= ChatPromptTemplate.from_messages(
[
("system","You are a helpful assistant."),
("human","{input}"),
]
)

parser= StrOutputParser()

chain= prompt| chat| parser

try:
    chain_response= chain.invoke(
{"input":"Explain the concept of serverless computing in simple terms."}
)
print("--- Chain Response ---")
print(chain_response)
except Exceptionas e:
print(f"Error running chain:{e}")


# Async chain
try:
    async_chain_response=await chain.ainvoke(
{"input":"What are the benefits of using RunPod for AI/ML workloads?"}
)
print("--- Async Chain Response ---")
print(async_chain_response)
except Exceptionas e:
print(f"Error running async chain:{e}")