66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
from typing import Any, Dict, List
|
|
|
|
import requests
|
|
|
|
# Required: export DATABRICKS_TOKEN with a valid Databricks PAT before running.
|
|
DATABRICKS_TOKEN = "dapied25006cd5ec824b138dfa5e82e2630c" #User1
|
|
|
|
if not DATABRICKS_TOKEN:
|
|
raise RuntimeError("Set DATABRICKS_TOKEN with a valid Databricks PAT.")
|
|
|
|
# Optional: override DATABRICKS_BASE_URL to point at your workspace host.
|
|
BASE_URL = os.environ.get(
|
|
"DATABRICKS_BASE_URL",
|
|
"https://adb-7405604858124224.4.azuredatabricks.net",
|
|
)
|
|
ENDPOINT = f"{BASE_URL.rstrip('/')}/serving-endpoints/anthropic/v1/messages"
|
|
# Adjust DEFAULT_MODEL if your workspace exposes a different deployed model name.
|
|
DEFAULT_MODEL = "databricks-claude-sonnet-4-5"
|
|
|
|
|
|
def build_payload(prompt: str, model: str, max_tokens: int) -> Dict[str, Any]:
|
|
return {
|
|
"model": model,
|
|
"max_tokens": max_tokens,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
}
|
|
|
|
|
|
def call_message_api(
|
|
prompt: str, model: str = DEFAULT_MODEL, max_tokens: int = 500
|
|
) -> Dict[str, Any]:
|
|
payload = build_payload(prompt, model, max_tokens)
|
|
headers = {
|
|
"Authorization": f"Bearer {DATABRICKS_TOKEN}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
response = requests.post(ENDPOINT, headers=headers, json=payload, timeout=30)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
def extract_text(content: Any) -> str:
|
|
if isinstance(content, list):
|
|
text_blocks: List[str] = []
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") == "text":
|
|
text_blocks.append(str(block.get("text", "")))
|
|
if text_blocks:
|
|
return "\n".join(text_blocks)
|
|
return str(content)
|
|
|
|
|
|
def main(argv: List[str]) -> None:
|
|
prompt = " ".join(argv) if argv else "What is an LLM agent?"
|
|
result = call_message_api(prompt)
|
|
content = result.get("content")
|
|
output = extract_text(content)
|
|
print(output or json.dumps(result, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main(sys.argv[1:])
|