#!/usr/bin/env python3
"""
WhatsApp Message Scheduler - Send scheduled messages via WhatsApp Web
"""

import argparse
import asyncio
import json
import logging
import os
import random
import re
from datetime import datetime, timedelta
from pathlib import Path

def setup_environment(debug: bool):
	if not debug:
		os.environ['BROWSER_USE_SETUP_LOGGING'] = 'false'
		os.environ['BROWSER_USE_LOGGING_LEVEL'] = 'critical'
		logging.getLogger().setLevel(logging.CRITICAL)
	else:
		os.environ['BROWSER_USE_SETUP_LOGGING'] = 'true'
		os.environ['BROWSER_USE_LOGGING_LEVEL'] = 'info'

parser = argparse.ArgumentParser(description='WhatsApp Scheduler - Send scheduled messages via WhatsApp Web')
parser.add_argument('--debug', action='store_true', help='Debug mode: show browser and verbose logs')
parser.add_argument('--test', action='store_true', help='Test mode: show what messages would be sent without sending them')
parser.add_argument('--auto', action='store_true', help='Auto mode: respond to unread messages every 30 minutes')
args = parser.parse_args()
setup_environment(args.debug)

from browser_use import Agent, BrowserSession
from browser_use.llm.google import ChatGoogle

GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY') or os.getenv('GEMINI_API_KEY')

USER_DATA_DIR = Path.home() / '.config' / 'whatsapp_scheduler' / 'browser_profile'
USER_DATA_DIR.mkdir(parents=True, exist_ok=True)
STORAGE_STATE_FILE = USER_DATA_DIR / 'storage_state.json'

async def parse_messages():
	"""Parse messages.txt and extract scheduling info"""
	messages_file = Path('messages.txt')
	if not messages_file.exists():
		print('❌ messages.txt not found!')
		return []

	import aiofiles

	async with aiofiles.open(messages_file) as f:
		content = await f.read()

	llm = ChatGoogle(model='gemini-2.0-flash-exp', temperature=0.1, api_key=GOOGLE_API_KEY)

	now = datetime.now()
	prompt = f"""
	Parse these WhatsApp message instructions and extract:
	1. Contact name (extract just the name, not descriptions)
	2. Message content (what to send)
	3. Date and time (when to send)

	Current date/time: {now.strftime('%Y-%m-%d %H:%M')}
	Today is: {now.strftime('%Y-%m-%d')}
	Current time is: {now.strftime('%H:%M')}

	Instructions:
	{content}

	Return ONLY a JSON array with format:
	[{{"contact": "name", "message": "text", "datetime": "YYYY-MM-DD HH:MM"}}]

	CRITICAL: Transform instructions into actual messages:

	QUOTED TEXT → Use exactly as-is:
	- Text in "quotes" becomes the exact message

	UNQUOTED INSTRUCTIONS → Generate actual content:
	- If it's an instruction to write something → write the actual thing
	- If it's an instruction to tell someone something → write what to tell them
	- If it's an instruction to remind someone → write the actual reminder
	- For multi-line content like poems: use single line with spacing, not line breaks

	DO NOT copy the instruction - create the actual message content!

	Time Rules:
	- If only time given (like "at 15:30"), use TODAY
	- If no date specified, assume TODAY
	- If no year given, use current year
	- Default time is 9:00 if not specified
	- Extract names from parentheses: "hinge date (Camila)" → "Camila"
	- "tomorrow" means {(now + timedelta(days=1)).strftime('%Y-%m-%d')}
	- "next tuesday" or similar means the next occurrence of that day
	"""

	from browser_use.llm.messages import UserMessage

	response = await llm.ainvoke([UserMessage(content=prompt)])
	response_text = response.completion if hasattr(response, 'completion') else str(response)

	# Extract JSON
	json_match = re.search(r'\[.*?\]', response_text, re.DOTALL)
	if json_match:
		try:
			messages = json.loads(json_match.group())
			for msg in messages:
				if 'message' in msg:
					msg['message'] = re.sub(r'\n+', ' • ', msg['message'])
					msg['message'] = re.sub(r'\s+', ' ', msg['message']).strip()
			return messages
		except json.JSONDecodeError:
			pass
	return []

async def send_message(contact, message):
	"""Send a WhatsApp message"""
	print(f'\n📱 Sending to {contact}: {message}')

	llm = ChatGoogle(model='gemini-2.0-flash-exp', temperature=0.3, api_key=GOOGLE_API_KEY)

	task = f"""
	Send WhatsApp message:
	1. Go to https://web.whatsapp.com
	2. Search for contact: {contact}
	3. Click on the contact
	4. Type message: {message}
	5. Press Enter to send
	6. Confirm sent
	"""

	browser = BrowserSession(
		headless=not args.debug,  # headless=False only when debug=True
		user_data_dir=str(USER_DATA_DIR),
		storage_state=str(STORAGE_STATE_FILE) if STORAGE_STATE_FILE.exists() else None,
	)

	agent = Agent(task=task, llm=llm, browser_session=browser)
	await agent.run()
	print(f'✅ Sent to {contact}')

async def auto_respond_to_unread():
	"""Click unread tab and respond to messages"""
	print('\nAuto-responding to unread messages...')

	llm = ChatGoogle(model='gemini-2.0-flash-exp', temperature=0.3, api_key=GOOGLE_API_KEY)

	task = """
	1. Go to https://web.whatsapp.com
	2. Wait for page to load
	3. Click on the "Unread" filter tab
	4. If there are unread messages:
	   - Click on each unread chat
	   - Read the last message
	   - Generate and send a friendly, contextual response
	   - Move to next unread chat
	5. Report how many messages were responded to
	"""

	browser = BrowserSession(
		headless=not args.debug,
		user_data_dir=str(USER_DATA_DIR),
		storage_state=str(STORAGE_STATE_FILE) if STORAGE_STATE_FILE.exists() else None,
	)

	agent = Agent(task=task, llm=llm, browser_session=browser)
	result = await agent.run()
	print('✅ Auto-response complete')
	return result

async def main():
	if not GOOGLE_API_KEY:
		print('❌ Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable')
		return

	print('WhatsApp Scheduler')
	print(f'Profile: {USER_DATA_DIR}')
	print()

	# Auto mode - respond to unread messages periodically
	if args.auto:
		print('AUTO MODE - Responding to unread messages every ~30 minutes')
		print('Press Ctrl+C to stop.\n')

		while True:
			try:
				await auto_respond_to_unread()

				# Wait 30 minutes +/- 5 minutes randomly
				wait_minutes = 30 + random.randint(-5, 5)
				print(f'\n⏰ Next check in {wait_minutes} minutes...')
				await asyncio.sleep(wait_minutes * 60)

			except KeyboardInterrupt:
				print('\n\nAuto mode stopped by user')
				break
			except Exception as e:
				print(f'\n❌ Error in auto mode: {e}')
				print('Waiting 5 minutes before retry...')
				await asyncio.sleep(300)
		return

	# Parse messages
	print('Parsing messages.txt...')
	messages = await parse_messages()

	if not messages:
		print('No messages found')
		return

	print(f'\nFound {len(messages)} messages:')
	for msg in messages:
		print(f'  • {msg["datetime"]}: {msg["message"][:30]}... to {msg["contact"]}')

	now = datetime.now()
	immediate = []
	future = []

	for msg in messages:
		msg_time = datetime.strptime(msg['datetime'], '%Y-%m-%d %H:%M')
		if msg_time <= now:
			immediate.append(msg)
		else:
			future.append(msg)

	if args.test:
		print('\n=== TEST MODE - Preview ===')
		if immediate:
			print(f'\nWould send {len(immediate)} past-due messages NOW:')
			for msg in immediate:
				print(f'  📱 To {msg["contact"]}: {msg["message"]}')
		if future:
			print(f'\nWould monitor {len(future)} future messages:')
			for msg in future:
				print(f'  ⏰ {msg["datetime"]}: To {msg["contact"]}: {msg["message"]}')
		print('\nTest mode complete. No messages sent.')
		return

	if immediate:
		print(f'\nSending {len(immediate)} past-due messages NOW...')
		for msg in immediate:
			await send_message(msg['contact'], msg['message'])

	if future:
		print(f'\n⏰ Monitoring {len(future)} future messages...')
		print('Press Ctrl+C to stop.\n')

		last_status = None

		while future:
			now = datetime.now()
			due = []
			remaining = []

			for msg in future:
				msg_time = datetime.strptime(msg['datetime'], '%Y-%m-%d %H:%M')
				if msg_time <= now:
					due.append(msg)
				else:
					remaining.append(msg)

			for msg in due:
				print(f'\n⏰ Time reached for {msg["contact"]}')
				await send_message(msg['contact'], msg['message'])

			future = remaining

			if future:
				next_msg = min(future, key=lambda x: datetime.strptime(x['datetime'], '%Y-%m-%d %H:%M'))
				current_status = f'Next: {next_msg["datetime"]} to {next_msg["contact"]}'

				if current_status != last_status:
					print(current_status)
					last_status = current_status

				await asyncio.sleep(30)  # Check every 30 seconds

	print('\n✅ All messages processed!')

if __name__ == '__main__':
	asyncio.run(main())


// Code generated by templ - DO NOT EDIT.

// templ: version: v0.3.1020
package metarefresh

//lint:file-ignore SA4006 This context is only used if a nested component is present.

import "github.com/a-h/templ"
import templruntime "github.com/a-h/templ/runtime"

import (
	"fmt"

	"github.com/TecharoHQ/anubis"
	"github.com/TecharoHQ/anubis/lib/localization"
)

func page(redir string, difficulty int, showMeta bool, loc *localization.SimpleLocalizer) templ.Component {
	return templruntime.GeneratedTemplate(func(templ_7745c5c3_Input templruntime.GeneratedComponentInput) (templ_7745c5c3_Err error) {
		templ_7745c5c3_W, ctx := templ_7745c5c3_Input.Writer, templ_7745c5c3_Input.Context
		if templ_7745c5c3_CtxErr := ctx.Err(); templ_7745c5c3_CtxErr != nil {
			return templ_7745c5c3_CtxErr
		}
		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templruntime.GetBuffer(templ_7745c5c3_W)
		if !templ_7745c5c3_IsBuffer {
			defer func() {
				templ_7745c5c3_BufErr := templruntime.ReleaseBuffer(templ_7745c5c3_Buffer)
				if templ_7745c5c3_Err == nil {
					templ_7745c5c3_Err = templ_7745c5c3_BufErr
				}
			}()
		}
		ctx = templ.InitializeContext(ctx)
		templ_7745c5c3_Var1 := templ.GetChildren(ctx)
		if templ_7745c5c3_Var1 == nil {
			templ_7745c5c3_Var1 = templ.NopComponent
		}
		ctx = templ.ClearChildren(ctx)
		templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 1, "
<img id=\"image\" style=\"width:100%;max-width:256px;\" src=\"")
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		var templ_7745c5c3_Var2 string
		templ_7745c5c3_Var2, templ_7745c5c3_Err = templ.ResolveAttributeValue(anubis.BasePrefix + "/.within.website/x/cmd/anubis/static/img/pensive.webp?cacheBuster=" + anubis.Version)
		if templ_7745c5c3_Err != nil {
			return templ.Error{Err: templ_7745c5c3_Err, FileName: `metarefresh.templ`, Line: 12, Col: 165}
		}
		_, templ_7745c5c3_Err = templ_7745c5c3_Buffer.WriteString(templ_7745c5c3_Var2)
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 2, "\"> <img style=\"display:none;\" style=\"width:100%;max-width:256px;\" src=\"")
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		var templ_7745c5c3_Var3 string
		templ_7745c5c3_Var3, templ_7745c5c3_Err = templ.ResolveAttributeValue(anubis.BasePrefix + "/.within.website/x/cmd/anubis/static/img/happy.webp?cacheBuster=" + anubis.Version)
		if templ_7745c5c3_Err != nil {
			return templ.Error{Err: templ_7745c5c3_Err, FileName: `metarefresh.templ`, Line: 13, Col: 174}
		}
		_, templ_7745c5c3_Err = templ_7745c5c3_Buffer.WriteString(templ_7745c5c3_Var3)
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 3, "\">
")
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		var templ_7745c5c3_Var4 string
		templ_7745c5c3_Var4, templ_7745c5c3_Err = templ.JoinStringErrs(loc.T("loading"))
		if templ_7745c5c3_Err != nil {
			return templ.Error{Err: templ_7745c5c3_Err, FileName: `metarefresh.templ`, Line: 14, Col: 35}
		}
		_, templ_7745c5c3_Err = templ_7745c5c3_Buffer.WriteString(templ.EscapeString(templ_7745c5c3_Var4))
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 4, "

")
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		var templ_7745c5c3_Var5 string
		templ_7745c5c3_Var5, templ_7745c5c3_Err = templ.JoinStringErrs(loc.T("connection_security"))
		if templ_7745c5c3_Err != nil {
			return templ.Error{Err: templ_7745c5c3_Err, FileName: `metarefresh.templ`, Line: 15, Col: 35}
		}
		_, templ_7745c5c3_Err = templ_7745c5c3_Buffer.WriteString(templ.EscapeString(templ_7745c5c3_Var5))
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 5, "
")
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		if showMeta {
			templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 6, "<meta http-equiv=\"refresh\" content=\"")
			if templ_7745c5c3_Err != nil {
				return templ_7745c5c3_Err
			}
			var templ_7745c5c3_Var6 string
			templ_7745c5c3_Var6, templ_7745c5c3_Err = templ.ResolveAttributeValue(fmt.Sprintf("%d; url=%s", difficulty+1, redir))
			if templ_7745c5c3_Err != nil {
				return templ.Error{Err: templ_7745c5c3_Err, FileName: `metarefresh.templ`, Line: 17, Col: 86}
			}
			_, templ_7745c5c3_Err = templ_7745c5c3_Buffer.WriteString(templ_7745c5c3_Var6)
			if templ_7745c5c3_Err != nil {
				return templ_7745c5c3_Err
			}
			templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 7, "\">")
			if templ_7745c5c3_Err != nil {
				return templ_7745c5c3_Err
			}
		}
		templ_7745c5c3_Err = templruntime.WriteString(templ_7745c5c3_Buffer, 8, "
")
		if templ_7745c5c3_Err != nil {
			return templ_7745c5c3_Err
		}
		return nil
	})
}

var _ = templruntime.GeneratedTemplate


import { describe, it, expect, vi } from 'vitest';

const { __test__ } = await import('./search.js');
const command = __test__.command;

function createPageMock(evaluateResult = []) {
  return {
    goto: vi.fn().mockResolvedValue(undefined),
    wait: vi.fn().mockResolvedValue(undefined),
    evaluate: vi.fn().mockResolvedValue(evaluateResult),
  };
}

describe('yahoo search', () => {
  it('should register as a valid command', () => {
    expect(command).toBeDefined();
    expect(command.site).toBe('yahoo');
    expect(command.name).toBe('search');
    expect(command.access).toBe('read');
    expect(command.browser).toBe(true);
    expect(command.strategy).toBe('public');
    expect(command.domain).toBe('search.yahoo.com');
  });

  it('should define keyword positional arg', () => {
    const kwArg = command.args.find(a => a.name === 'keyword');
    expect(kwArg).toBeDefined();
    expect(kwArg.positional).toBe(true);
    expect(kwArg.required).toBe(true);
  });

  it('should define limit arg with default 7', () => {
    const limitArg = command.args.find(a => a.name === 'limit');
    expect(limitArg).toBeDefined();
    expect(limitArg.type).toBe('int');
    expect(limitArg.default).toBe(7);
  });

  it('should define output columns', () => {
    expect(command.columns).toContain('rank');
    expect(command.columns).toContain('title');
    expect(command.columns).toContain('url');
    expect(command.columns).toContain('snippet');
  });

  it('rejects empty query, invalid limit, and invalid page before navigation', async () => {
    const page = createPageMock();
    await expect(command.func(page, { keyword: ' ', limit: 5 })).rejects.toMatchObject({ code: 'ARGUMENT' });
    await expect(command.func(page, { keyword: 'opencli', limit: 8 })).rejects.toMatchObject({ code: 'ARGUMENT' });
    await expect(command.func(page, { keyword: 'opencli', limit: 5, page: 0 })).rejects.toMatchObject({ code: 'ARGUMENT' });
    expect(page.goto).not.toHaveBeenCalled();
  });

  it('decodes Yahoo redirect URLs and assigns listing rank', async () => {
    const page = createPageMock({
      session: 'site:yahoo',
      data: [[
        'OpenCLI',
        'https://r.search.yahoo.com/_ylt=x/RU=https%3A%2F%2Fgithub.com%2Fjackwener%2FOpenCLI/RK=2/RS=x',
        'CLI browser tooling',
      ]],
    });

    await expect(command.func(page, { keyword: 'opencli', limit: 1, page: 2 })).resolves.toEqual([{
      rank: 8,
      title: 'OpenCLI',
      url: 'https://github.com/jackwener/OpenCLI',
      snippet: 'CLI browser tooling',
    }]);
  });

  it('drops decoded Yahoo redirect targets that are not http(s) URLs', async () => {
    const page = createPageMock([
      [
        'Bad redirect',
        'https://r.search.yahoo.com/_ylt=x/RU=javascript%3Aalert(1)/RK=2/RS=x',
        'should not be emitted',
      ],
    ]);

    await expect(command.func(page, { keyword: 'opencli', limit: 1 })).rejects.toMatchObject({
      code: 'EMPTY_RESULT',
    });
  });

  it('fails typed instead of silently returning [] for malformed extraction payloads', async () => {
    const page = createPageMock({ rows: [] });

    await expect(command.func(page, { keyword: 'opencli', limit: 1 })).rejects.toMatchObject({
      code: 'COMMAND_EXEC',
      message: expect.stringContaining('payload shape'),
    });
  });
});


import re
import base64
import binascii
from urllib.parse import parse_qs, urlparse
from typing import Any, BinaryIO
from bs4 import BeautifulSoup

from .._base_converter import DocumentConverter, DocumentConverterResult
from .._stream_info import StreamInfo
from ._markdownify import _CustomMarkdownify

ACCEPTED_MIME_TYPE_PREFIXES = [
    "text/html",
    "application/xhtml",
]

ACCEPTED_FILE_EXTENSIONS = [
    ".html",
    ".htm",
]

class BingSerpConverter(DocumentConverter):
    """
    Handle Bing results pages (only the organic search results).
    NOTE: It is better to use the Bing API
    """

    def accepts(
        self,
        file_stream: BinaryIO,
        stream_info: StreamInfo,
        **kwargs: Any,  # Options to pass to the converter
    ) -> bool:
        """
        Make sure we're dealing with HTML content *from* Bing.
        """

        url = stream_info.url or ""
        mimetype = (stream_info.mimetype or "").lower()
        extension = (stream_info.extension or "").lower()

        if not re.search(r"^https://www\.bing\.com/search\?q=", url):
            # Not a Bing SERP URL
            return False

        if extension in ACCEPTED_FILE_EXTENSIONS:
            return True

        for prefix in ACCEPTED_MIME_TYPE_PREFIXES:
            if mimetype.startswith(prefix):
                return True

        # Not HTML content
        return False

    def convert(
        self,
        file_stream: BinaryIO,
        stream_info: StreamInfo,
        **kwargs: Any,  # Options to pass to the converter
    ) -> DocumentConverterResult:
        assert stream_info.url is not None

        # Parse the query parameters
        parsed_params = parse_qs(urlparse(stream_info.url).query)
        query = parsed_params.get("q", [""])[0]

        # Parse the stream
        encoding = "utf-8" if stream_info.charset is None else stream_info.charset
        soup = BeautifulSoup(file_stream, "html.parser", from_encoding=encoding)

        # Clean up some formatting
        for tptt in soup.find_all(class_="tptt"):
            if hasattr(tptt, "string") and tptt.string:
                tptt.string += " "
        for slug in soup.find_all(class_="algoSlug_icon"):
            slug.extract()

        # Parse the algorithmic results
        _markdownify = _CustomMarkdownify(**kwargs)
        results = list()
        for result in soup.find_all(class_="b_algo"):
            if not hasattr(result, "find_all"):
                continue

            # Rewrite redirect urls
            for a in result.find_all("a", href=True):
                parsed_href = urlparse(a["href"])
                qs = parse_qs(parsed_href.query)

                # The destination is contained in the u parameter,
                # but appears to be base64 encoded, with some prefix
                if "u" in qs:
                    u = (
                        qs["u"][0][2:].strip() + "=="
                    )  # Python 3 doesn't care about extra padding

                    try:
                        # RFC 4648 / Base64URL variant, which uses "-" and "_"
                        a["href"] = base64.b64decode(u, altchars="-_").decode("utf-8")
                    except UnicodeDecodeError:
                        pass
                    except binascii.Error:
                        pass

            # Convert to markdown
            md_result = _markdownify.convert_soup(result).strip()
            lines = [line.strip() for line in re.split(r"\n+", md_result)]
            results.append("\n".join([line for line in lines if len(line) > 0]))

        webpage_text = (
            f"## A Bing search for '{query}' found the following results:\n\n"
            + "\n\n".join(results)
        )

        return DocumentConverterResult(
            markdown=webpage_text,
            title=None if soup.title is None else soup.title.string,
        )


package io.unitycatalog.server.utils;

import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertThrows;

import io.unitycatalog.client.ApiClient;
import io.unitycatalog.client.ApiClientBuilder;
import io.unitycatalog.client.ApiException;
import io.unitycatalog.client.auth.TokenProvider;
import io.unitycatalog.client.delta.DeltaApiException;
import io.unitycatalog.client.delta.model.DeltaErrorType;
import io.unitycatalog.client.retry.JitterDelayRetryPolicy;
import io.unitycatalog.server.base.ServerConfig;
import io.unitycatalog.server.exception.ErrorCode;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import org.junit.jupiter.api.function.Executable;

public class TestUtils {
  public static final String CATALOG_NAME = "uc_testcatalog";
  public static final String SCHEMA_NAME = "uc_testschema";
  public static final String CATALOG_NAME2 = "uc_testcatalog2";
  public static final String SCHEMA_NAME2 = "uc_testschema2";
  public static final String TABLE_NAME = "uc_testtable";
  public static final String VOLUME_NAME = "uc_testvolume";
  public static final String FUNCTION_NAME = "uc_testfunction";
  public static final String MODEL_NAME = "uc_testmodel";
  public static final String MODEL_NEW_NAME = "uc_newtestmodel";
  public static final String SCHEMA_FULL_NAME = CATALOG_NAME + "." + SCHEMA_NAME;
  public static final String SCHEMA_NEW_NAME = "uc_newtestschema";
  public static final String SCHEMA_NEW_FULL_NAME = CATALOG_NAME + "." + SCHEMA_NEW_NAME;
  public static final String SCHEMA_NEW_COMMENT = "new test comment";
  public static final String TABLE_FULL_NAME = CATALOG_NAME + "." + SCHEMA_NAME + "." + TABLE_NAME;
  public static final String VOLUME_FULL_NAME =
      CATALOG_NAME + "." + SCHEMA_NAME + "." + VOLUME_NAME;
  public static final String FUNCTION_FULL_NAME =
      CATALOG_NAME + "." + SCHEMA_NAME + "." + FUNCTION_NAME;
  public static final String MODEL_FULL_NAME = CATALOG_NAME + "." + SCHEMA_NAME + "." + MODEL_NAME;
  public static final String MODEL_NEW_FULL_NAME =
      CATALOG_NAME + "." + SCHEMA_NAME + "." + MODEL_NEW_NAME;
  public static final String COMMENT = "test comment";
  public static final String COMMENT2 = "test comment 2";
  public static final String CATALOG_NEW_NAME = "uc_newtestcatalog";
  public static final String CATALOG_NEW_COMMENT = "new test comment";
  public static final String MODEL_NEW_COMMENT = "new test model comment";
  public static final String VOLUME_NEW_NAME = "uc_newtestvolume";
  public static final String VOLUME_NEW_FULL_NAME =
      CATALOG_NAME + "." + SCHEMA_NAME + "." + VOLUME_NEW_NAME;
  public static final String MV_COMMENT = "model version comment";
  public static final String MV_SOURCE = "model version source";
  public static final String MV_RUNID = "model version runId";
  public static final String MV_SOURCE2 = "model version source 2";
  public static final String MV_RUNID2 = "model version runId 2";
  public static final String TEST_AWS_MASTER_ROLE_ARN =
      "arn:aws:iam::1234567:role/UCMasterRole-EXAMPLE";
  public static final String TEST_AWS_MASTER_ROLE_ACCESS_KEY = "masterRoleAccessKey";
  public static final String TEST_AWS_MASTER_ROLE_SECRET_KEY = "masterRoleSecretKey";
  public static final String TEST_AWS_REGION = "us-west-2";

  public static final Map<String, String> PROPERTIES =
      new HashMap<>(Map.of("prop1", "value1", "prop2", "value2"));
  public static final Map<String, String> NEW_PROPERTIES =
      new HashMap<>(Map.of("prop2", "value22", "prop3", "value33"));
  public static final String COMMON_ENTITY_NAME = "zz_uc_common_entity_name";

  public static ApiClient createApiClient(ServerConfig serverConfig) {
    URI uri = URI.create(serverConfig.getServerUrl());
    String token = serverConfig.getAuthToken() != null ? serverConfig.getAuthToken() : "";
    return ApiClientBuilder.create()
        .uri(uri)
        .tokenProvider(TokenProvider.create(Map.of("type", "static", "token", token)))
        .retryPolicy(JitterDelayRetryPolicy.builder().maxAttempts(1).build())
        .build();
  }

  public static void assertApiException(
      Executable executable, ErrorCode errorCode, String containsMessage) {
    ApiException ex = assertThrows(ApiException.class, executable);
    // Check the message first. When tests fail due to mismatching error, the message can tell us
    // more.
    assertThat(ex.getMessage()).contains(containsMessage);
    assertThat(ex.getCode()).isEqualTo(errorCode.getHttpStatus().code());
  }

  /**
   * Asserts the call fails by checking the HTTP status code is {@code expectedStatus}. Use for
   * body-less responses; otherwise use {@link #assertApiException} / {@link
   * #assertDeltaApiException}.
   */
  public static void assertApiExceptionStatusOnly(Executable executable, int expectedStatus) {
    ApiException ex = assertThrows(ApiException.class, executable);
    assertThat(ex.getCode()).isEqualTo(expectedStatus);
  }

  public static void assertDeltaApiException(
      Executable executable, DeltaErrorType expectedType, String expectedMessageSubstring) {
    int expectedCode = ErrorCode.getDeltaHttpStatus(expectedType.getValue()).code();
    ApiException ex = assertThrows(ApiException.class, executable);
    // Check message first for better diagnostics on failure (includes the full response body)
    assertThat(ex.getMessage()).contains(expectedMessageSubstring);
    assertThat(ex.getCode()).isEqualTo(expectedCode);
    Optional<DeltaApiException> deltaExOpt = DeltaApiException.from(ex);
    assertThat(deltaExOpt)
        .as("Failed to parse Delta error response: " + ex.getResponseBody())
        .isPresent();
    DeltaApiException delta = deltaExOpt.get();
    assertThat(delta.getErrorCode()).isEqualTo(expectedCode);
    assertThat(delta.getErrorType()).isEqualTo(expectedType);
    assertThat(delta.getErrorMessage()).contains(expectedMessageSubstring);
  }

  /**
   * Asserts the call fails with PERMISSION_DENIED (HTTP 403) and the exception message contains
   * {@code containsMessage}. Use this when the test cares that a specific authz check fired -- e.g.
   * a staging-table ownership check -- rather than just that something produced a 403.
   */
  public static void assertPermissionDenied(Executable executable, String containsMessage) {
    assertApiException(executable, ErrorCode.PERMISSION_DENIED, containsMessage);
  }

  /**
   * Asserts the call fails with PERMISSION_DENIED (HTTP 403). Only checks the error code and the
   * generic {@code "PERMISSION_DENIED"} marker in the message; use {@link
   * #assertPermissionDenied(Executable, String)} when the specific cause matters.
   */
  public static void assertPermissionDenied(Executable executable) {
    assertPermissionDenied(executable, "PERMISSION_DENIED");
  }

  /**
   * Raw-HTTP counterpart to {@link #assertApiException}. Use when the generated SDK can't reach the
   * failure mode (e.g. the SDK always serializes a body, so you can't exercise body-less
   * authorization paths with it).
   */
  public static void assertHttpApiException(
      HttpResponse<String> response, ErrorCode errorCode, String containsMessage) {
    // Check the body first. When tests fail due to mismatching error, the body can tell us more.
    assertThat(response.body()).contains(containsMessage).contains(errorCode.name());
    assertThat(response.statusCode()).isEqualTo(errorCode.getHttpStatus().code());
  }

  /**
   * Sends a body-less POST to the given path. The SDK always attaches a serialized body, so raw
   * HTTP is the only way to reach the body-less code path (used to exercise {@link
   * io.unitycatalog.server.auth.decorator.AuthorizationGateConverter}'s silent-skip denial).
   */
  public static HttpResponse<String> sendRawEmptyPost(ServerConfig config, String path)
      throws Exception {
    HttpRequest.Builder reqBuilder =
        HttpRequest.newBuilder()
            .uri(URI.create(config.getServerUrl() + path))
            .POST(HttpRequest.BodyPublishers.noBody());
    if (config.getAuthToken() != null && !config.getAuthToken().isEmpty()) {
      reqBuilder.header("Authorization", "Bearer " + config.getAuthToken());
    }
    return HttpClient.newHttpClient()
        .send(reqBuilder.build(), HttpResponse.BodyHandlers.ofString());
  }
}


# Reference Card

> **Quick-reference for all 14+ correct-by-construction patterns** with selection flowchart, pattern catalogue, composition rules, crate mapping, and types-as-guarantees cheat sheet.
>
> **Cross-references:** Every chapter — this is the lookup table for the entire book.

## Quick Reference: Correct-by-Construction Patterns

### Pattern Selection Guide

```text
Is the bug catastrophic if missed?
├── Yes → Can it be encoded in types?
│         ├── Yes → USE CORRECT-BY-CONSTRUCTION
│         └── No  → Runtime check + extensive testing
└── No  → Runtime check is fine
```

### Pattern Catalogue

| # | Pattern | Key Trait/Type | Prevents | Runtime Cost | Chapter |
|---|---------|---------------|----------|:------:|---------|
| 1 | Typed Commands | `trait IpmiCmd { type Response; }` | Wrong response type | Zero | ch02 |
| 2 | Single-Use Types | `struct Nonce` (not Clone/Copy) | Nonce/key reuse | Zero | ch03 |
| 3 | Capability Tokens | `struct AdminToken { _private: () }` | Unauthorised access | Zero | ch04 |
| 4 | Type-State | `Session<Active>` | Protocol violations | Zero | ch05 |
| 5 | Dimensional Types | `struct Celsius(f64)` | Unit confusion | Zero | ch06 |
| 6 | Validated Boundaries | `struct ValidFru` (via TryFrom) | Unvalidated data use | Parse once | ch07 |
| 7 | Capability Mixins | `trait FanDiagMixin: HasSpi + HasI2c` | Missing bus access | Zero | ch08 |
| 8 | Phantom Types | `Register<Width16>` | Width/direction mismatch | Zero | ch09 |
| 9 | Sentinel → Option | `Option<u8>` (not `0xFF`) | Sentinel-as-value bugs | Zero | ch11 |
| 10 | Sealed Traits | `trait Cmd: private::Sealed` | Unsound external impls | Zero | ch11 |
| 11 | Non-Exhaustive Enums | `#[non_exhaustive] enum Sku` | Silent match fallthrough | Zero | ch11 |
| 12 | Typestate Builder | `DerBuilder<Set, Missing>` | Incomplete construction | Zero | ch11 |
| 13 | FromStr Validation | `impl FromStr for DiagLevel` | Unvalidated string input | Parse once | ch11 |
| 14 | Const-Generic Size | `RegisterBank<const N: usize>` | Buffer size mismatch | Zero | ch11 |
| 15 | Safe `unsafe` Wrapper | `MmioRegion::read_u32()` | Unchecked MMIO/FFI | Zero | ch11 |
| 16 | Async Type-State | `AsyncSession<Active>` | Async protocol violations | Zero | ch11 |
| 17 | Const Assertions | `SdrSensorId<const N: u8>` | Invalid compile-time IDs | Zero | ch11 |
| 18 | Session Types | `Chan<SendRequest>` | Out-of-order channel ops | Zero | ch11 |
| 19 | Pin Self-Referential | `Pin<Box<StreamParser>>` | Dangling intra-struct pointer | Zero | ch11 |
| 20 | RAII / Drop | `impl Drop for Session` | Resource leak on any exit path | Zero | ch11 |
| 21 | Error Type Hierarchy | `#[derive(Error)] enum DiagError` | Silent error swallowing | Zero | ch11 |
| 22 | `#[must_use]` | `#[must_use] struct Token` | Silently dropped values | Zero | ch11 |

### Composition Rules

```text
Capability Token + Type-State = Authorised state transitions
Typed Command + Dimensional Type = Physically-typed responses
Validated Boundary + Phantom Type = Typed register access on validated config
Capability Mixin + Typed Command = Bus-aware typed operations
Single-Use Type + Type-State = Consume-on-transition protocols
Sealed Trait + Typed Command = Closed, sound command set
Sentinel → Option + Validated Boundary = Clean parse-once pipeline
Typestate Builder + Capability Token = Proof-of-complete construction
FromStr + #[non_exhaustive] = Evolvable, fail-fast enum parsing
Const-Generic Size + Validated Boundary = Sized, validated protocol buffers
Safe unsafe Wrapper + Phantom Type = Typed, safe MMIO access
Async Type-State + Capability Token = Authorised async transitions
Session Types + Typed Command = Fully-typed request-response channels
Pin + Type-State = Self-referential state machines that can't move
RAII (Drop) + Type-State = State-dependent cleanup guarantees
Error Hierarchy + Validated Boundary = Typed parse errors with exhaustive handling
#[must_use] + Single-Use Type = Hard-to-ignore, hard-to-reuse tokens
```

### Anti-Patterns to Avoid

| Anti-Pattern | Why It's Wrong | Correct Alternative |
|-------------|---------------|-------------------|
| `fn read_sensor() -> f64` | Unitless — could be °C, °F, or RPM | `fn read_sensor() -> Celsius` |
| `fn encrypt(nonce: &[u8; 12])` | Nonce can be reused (borrow) | `fn encrypt(nonce: Nonce)` (move) |
| `fn admin_op(is_admin: bool)` | Caller can lie (`true`) | `fn admin_op(_: &AdminToken)` |
| `fn send(session: &Session)` | No state guarantee | `fn send(session: &Session<Active>)` |
| `fn process(data: &[u8])` | Not validated | `fn process(data: &ValidFru)` |
| `Clone` on ephemeral keys | Defeats single-use guarantee | Don't derive Clone |
| `let vendor_id: u16 = 0xFFFF` | Sentinel carried internally | `let vendor_id: Option<u16> = None` |
| `fn route(level: &str)` with fallback | Typos silently default | `let level: DiagLevel = s.parse()?` |
| `Builder::new().finish()` without fields | Incomplete object constructed | Typestate builder: `finish()` gated on `Set` |
| `let buf: Vec<u8>` for fixed-size HW buffer | Size only checked at runtime | `RegisterBank<4096>` (const generic) |
| Raw `unsafe { ptr::read(...) }` scattered | UB risk, unauditable | `MmioRegion::read_u32()` safe wrapper |
| `async fn transition(&mut self)` | Mutable borrows don't enforce state | `async fn transition(self) -> NextState` |
| `fn cleanup()` called manually | Forgotten on early return / panic | `impl Drop` — compiler inserts call |
| `fn op() -> Result<T, String>` | Opaque error, no variant matching | `fn op() -> Result<T, DiagError>` enum |

### Mapping to a Diagnostics Codebase

| Module | Applicable Pattern(s) |
|---------------------|----------------------|
| `protocol_lib` | Typed commands, type-state sessions |
| `thermal_diag` | Capability mixins, dimensional types |
| `accel_diag` | Validated boundaries, phantom registers |
| `network_diag` | Type-state (link training), capability tokens |
| `pci_topology` | Phantom types (register width), validated config, sentinel → Option |
| `event_handler` | Single-use audit tokens, capability tokens, FromStr (Component) |
| `event_log` | Validated boundaries (SEL record parsing) |
| `compute_diag` | Dimensional types (temperature, frequency) |
| `memory_diag` | Validated boundaries (SPD data), dimensional types |
| `switch_diag` | Type-state (port enumeration), phantom types |
| `config_loader` | FromStr (DiagLevel, FaultStatus, DiagAction) |
| `log_analyzer` | Validated boundaries (CompiledPatterns) |
| `diag_framework` | Typestate builder (DerBuilder), session types (orchestrator↔worker) |
| `topology_lib` | Const-generic register banks, safe MMIO wrappers |

### Types as Guarantees — Quick Mapping

| Guarantee | Rust Equivalent | Example |
|-----------|----------------|---------|
| "This proof exists" | A type | `AdminToken` |
| "I have the proof" | A value of that type | `let tok = authenticate()?;` |
| "A implies B" | Function `fn(A) -> B` | `fn activate(AdminToken) -> Session<Active>` |
| "Both A and B" | Tuple `(A, B)` or multi-param | `fn op(a: &AdminToken, b: &LinkTrained)` |
| "Either A or B" | `enum { A(A), B(B) }` or `Result<A, B>` | `Result<Session<Active>, Error>` |
| "Always true" | `()` (unit type) | Always constructible |
| "Impossible" | `!` (never type) or `enum Void {}` | Can never be constructed |

---


"""Azure Content Understanding converter for MarkItDown.

Converts files using Azure Content Understanding (CU) for high-quality,
multi-modal extraction with structured field output. Supports documents,
images, audio, and video. Fields are serialized as YAML front matter via
the CU SDK's ``to_llm_input()`` helper.

Install dependencies: ``pip install 'markitdown[az-content-understanding]'``
"""

import sys
import os
from typing import BinaryIO, Any, List, Optional, Dict
from enum import Enum

from .._base_converter import DocumentConverter, DocumentConverterResult
from .._stream_info import StreamInfo
from .._exceptions import MissingDependencyException

# Try loading optional dependencies — save error for later
_dependency_exc_info = None
try:
    from azure.ai.contentunderstanding import ContentUnderstandingClient, to_llm_input
    from azure.core.credentials import AzureKeyCredential, TokenCredential
    from azure.core.pipeline.policies import UserAgentPolicy
    from azure.identity import DefaultAzureCredential
except ImportError:
    _dependency_exc_info = sys.exc_info()

    # Stub classes for type hinting
    class AzureKeyCredential:  # type: ignore[no-redef]
        pass

    class TokenCredential:  # type: ignore[no-redef]
        pass

    class ContentUnderstandingClient:  # type: ignore[no-redef]
        pass

    class UserAgentPolicy:  # type: ignore[no-redef]
        pass

    class DefaultAzureCredential:  # type: ignore[no-redef]
        pass

    def to_llm_input(*args, **kwargs):  # type: ignore[no-redef]
        pass

# ---------------------------------------------------------------------------
# File type enum and routing tables
# ---------------------------------------------------------------------------

class ContentUnderstandingFileType(str, Enum):
    """Supported file types for Content Understanding conversion."""

    # Documents
    PDF = "pdf"
    DOCX = "docx"
    PPTX = "pptx"
    XLSX = "xlsx"
    HTML = "html"
    TXT = "txt"
    MD = "md"
    RTF = "rtf"
    XML = "xml"

    # Email
    EML = "eml"
    MSG = "msg"

    # Images (document modality)
    JPEG = "jpeg"
    PNG = "png"
    BMP = "bmp"
    TIFF = "tiff"
    HEIF = "heif"

    # Video
    MP4 = "mp4"
    M4V = "m4v"
    MOV = "mov"
    AVI = "avi"
    MKV = "mkv"
    WEBM = "webm"
    FLV = "flv"
    WMV = "wmv"

    # Audio
    WAV = "wav"
    MP3 = "mp3"
    M4A = "m4a"
    FLAC = "flac"
    OGG = "ogg"
    AAC = "aac"
    WMA = "wma"

# Extension → file type
_EXTENSION_MAP: Dict[str, ContentUnderstandingFileType] = {
    # Documents
    ".pdf": ContentUnderstandingFileType.PDF,
    ".docx": ContentUnderstandingFileType.DOCX,
    ".pptx": ContentUnderstandingFileType.PPTX,
    ".xlsx": ContentUnderstandingFileType.XLSX,
    ".html": ContentUnderstandingFileType.HTML,
    ".txt": ContentUnderstandingFileType.TXT,
    ".md": ContentUnderstandingFileType.MD,
    ".rtf": ContentUnderstandingFileType.RTF,
    ".xml": ContentUnderstandingFileType.XML,
    # Email
    ".eml": ContentUnderstandingFileType.EML,
    ".msg": ContentUnderstandingFileType.MSG,
    # Images
    ".jpg": ContentUnderstandingFileType.JPEG,
    ".jpeg": ContentUnderstandingFileType.JPEG,
    ".jpe": ContentUnderstandingFileType.JPEG,
    ".png": ContentUnderstandingFileType.PNG,
    ".bmp": ContentUnderstandingFileType.BMP,
    ".tiff": ContentUnderstandingFileType.TIFF,
    ".heif": ContentUnderstandingFileType.HEIF,
    ".heic": ContentUnderstandingFileType.HEIF,
    # Video
    ".mp4": ContentUnderstandingFileType.MP4,
    ".m4v": ContentUnderstandingFileType.M4V,
    ".mov": ContentUnderstandingFileType.MOV,
    ".avi": ContentUnderstandingFileType.AVI,
    ".mkv": ContentUnderstandingFileType.MKV,
    ".webm": ContentUnderstandingFileType.WEBM,
    ".flv": ContentUnderstandingFileType.FLV,
    ".wmv": ContentUnderstandingFileType.WMV,
    # Audio
    ".wav": ContentUnderstandingFileType.WAV,
    ".mp3": ContentUnderstandingFileType.MP3,
    ".m4a": ContentUnderstandingFileType.M4A,
    ".flac": ContentUnderstandingFileType.FLAC,
    ".ogg": ContentUnderstandingFileType.OGG,
    ".aac": ContentUnderstandingFileType.AAC,
    ".wma": ContentUnderstandingFileType.WMA,
}

# MIME type prefixes for each file type
_MIME_PREFIXES: Dict[ContentUnderstandingFileType, List[str]] = {
    # Documents
    ContentUnderstandingFileType.PDF: ["application/pdf", "application/x-pdf"],
    ContentUnderstandingFileType.DOCX: [
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    ],
    ContentUnderstandingFileType.PPTX: [
        "application/vnd.openxmlformats-officedocument.presentationml"
    ],
    ContentUnderstandingFileType.XLSX: [
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    ],
    ContentUnderstandingFileType.HTML: ["text/html", "application/xhtml+xml"],
    ContentUnderstandingFileType.TXT: ["text/plain"],
    ContentUnderstandingFileType.MD: ["text/markdown"],
    ContentUnderstandingFileType.RTF: ["text/rtf", "application/rtf"],
    ContentUnderstandingFileType.XML: ["text/xml", "application/xml"],
    # Email
    ContentUnderstandingFileType.EML: ["message/rfc822"],
    ContentUnderstandingFileType.MSG: ["application/vnd.ms-outlook"],
    # Images
    ContentUnderstandingFileType.JPEG: ["image/jpeg"],
    ContentUnderstandingFileType.PNG: ["image/png"],
    ContentUnderstandingFileType.BMP: ["image/bmp"],
    ContentUnderstandingFileType.TIFF: ["image/tiff"],
    ContentUnderstandingFileType.HEIF: ["image/heif", "image/heic"],
    # Video
    ContentUnderstandingFileType.MP4: ["video/mp4"],
    ContentUnderstandingFileType.M4V: ["video/x-m4v"],
    ContentUnderstandingFileType.MOV: ["video/quicktime"],
    ContentUnderstandingFileType.AVI: ["video/x-msvideo"],
    ContentUnderstandingFileType.MKV: ["video/x-matroska"],
    ContentUnderstandingFileType.WEBM: ["video/webm"],
    ContentUnderstandingFileType.FLV: ["video/x-flv"],
    ContentUnderstandingFileType.WMV: ["video/x-ms-wmv"],
    # Audio
    ContentUnderstandingFileType.WAV: ["audio/wav", "audio/x-wav"],
    ContentUnderstandingFileType.MP3: ["audio/mpeg", "audio/mp3"],
    ContentUnderstandingFileType.M4A: ["audio/mp4", "audio/m4a", "audio/x-m4a"],
    ContentUnderstandingFileType.FLAC: ["audio/flac", "audio/x-flac"],
    ContentUnderstandingFileType.OGG: ["audio/ogg"],
    ContentUnderstandingFileType.AAC: ["audio/aac"],
    ContentUnderstandingFileType.WMA: ["audio/x-ms-wma"],
}

_MIME_ALIASES: Dict[str, str] = {
    "audio/x-wav": "audio/wav",
    "audio/x-flac": "audio/flac",
    "audio/x-m4a": "audio/mp4",
    "video/x-m4v": "video/mp4",
}

# File type → modality category
_DOCUMENT_TYPES = {
    ContentUnderstandingFileType.PDF,
    ContentUnderstandingFileType.DOCX,
    ContentUnderstandingFileType.PPTX,
    ContentUnderstandingFileType.XLSX,
    ContentUnderstandingFileType.HTML,
    ContentUnderstandingFileType.TXT,
    ContentUnderstandingFileType.MD,
    ContentUnderstandingFileType.RTF,
    ContentUnderstandingFileType.XML,
    ContentUnderstandingFileType.EML,
    ContentUnderstandingFileType.MSG,
}

_IMAGE_TYPES = {
    ContentUnderstandingFileType.JPEG,
    ContentUnderstandingFileType.PNG,
    ContentUnderstandingFileType.BMP,
    ContentUnderstandingFileType.TIFF,
    ContentUnderstandingFileType.HEIF,
}

_VIDEO_TYPES = {
    ContentUnderstandingFileType.MP4,
    ContentUnderstandingFileType.M4V,
    ContentUnderstandingFileType.MOV,
    ContentUnderstandingFileType.AVI,
    ContentUnderstandingFileType.MKV,
    ContentUnderstandingFileType.WEBM,
    ContentUnderstandingFileType.FLV,
    ContentUnderstandingFileType.WMV,
}

_AUDIO_TYPES = {
    ContentUnderstandingFileType.WAV,
    ContentUnderstandingFileType.MP3,
    ContentUnderstandingFileType.M4A,
    ContentUnderstandingFileType.FLAC,
    ContentUnderstandingFileType.OGG,
    ContentUnderstandingFileType.AAC,
    ContentUnderstandingFileType.WMA,
}

_PREBUILT_ANALYZERS = {
    "document": "prebuilt-documentSearch",
    "image": "prebuilt-documentSearch",
    "video": "prebuilt-videoSearch",
    "audio": "prebuilt-audioSearch",
}

# All supported file types (default set when file_types is None)
_ALL_FILE_TYPES = list(ContentUnderstandingFileType)

def _get_modality(file_type: ContentUnderstandingFileType) -> str:
    """Get the modality category for a file type."""
    if file_type in _DOCUMENT_TYPES:
        return "document"
    elif file_type in _IMAGE_TYPES:
        return "image"
    elif file_type in _VIDEO_TYPES:
        return "video"
    elif file_type in _AUDIO_TYPES:
        return "audio"
    raise ValueError(f"Unknown file type: {file_type}")

def _detect_file_type(
    stream_info: StreamInfo,
    file_types: Optional[List[ContentUnderstandingFileType]] = None,
) -> Optional[ContentUnderstandingFileType]:
    """Detect a supported CU file type from extension or MIME type."""
    allowed = set(file_types) if file_types is not None else None

    extension = (stream_info.extension or "").lower()
    file_type = _EXTENSION_MAP.get(extension)
    if file_type is not None and (allowed is None or file_type in allowed):
        return file_type

    mimetype = _clean_mime_type(stream_info.mimetype)
    if not mimetype:
        return None

    return _detect_file_type_from_mime(mimetype, allowed)

def _clean_mime_type(mimetype: Optional[str]) -> str:
    return (mimetype or "").split(";", 1)[0].strip().lower()

def _canonical_mime_type(mimetype: Optional[str]) -> str:
    cleaned = _clean_mime_type(mimetype)
    return _MIME_ALIASES.get(cleaned, cleaned) or "application/octet-stream"

def _content_type_for(
    file_type: ContentUnderstandingFileType,
    mimetype: Optional[str],
) -> str:
    """Resolve the content type to send to the CU API.

    Uses the resolved ``file_type`` as the source of truth so analyzer
    routing and payload metadata stay consistent. The caller-provided
    ``mimetype`` is only used when it is consistent with ``file_type``
    (e.g., to preserve subtype distinctions like ``image/heic`` vs
    ``image/heif``). When ``mimetype`` disagrees with the resolved
    ``file_type`` (e.g., ``.pdf`` extension with ``audio/mpeg``
    mimetype), the canonical MIME type for ``file_type`` is used.
    """
    prefixes = _MIME_PREFIXES.get(file_type, [])
    canonical = _canonical_mime_type(mimetype)

    # Use caller-provided MIME if it's consistent with the resolved file_type
    if prefixes and canonical != "application/octet-stream":
        for prefix in prefixes:
            if canonical.startswith(prefix):
                return canonical

    # Fallback: derive from the resolved file_type (single source of truth)
    if prefixes:
        return _canonical_mime_type(prefixes[0])

    return canonical

def _detect_file_type_from_mime(
    mimetype: str,
    allowed: Optional[set[ContentUnderstandingFileType]],
) -> Optional[ContentUnderstandingFileType]:
    for candidate, prefixes in _MIME_PREFIXES.items():
        if allowed is not None and candidate not in allowed:
            continue
        for prefix in prefixes:
            if mimetype.startswith(prefix):
                return candidate
    return None

# ---------------------------------------------------------------------------
# Smart routing: base_analyzer_id → modality mapping
# ---------------------------------------------------------------------------

_BASE_TO_MODALITY: Dict[str, str] = {
    "prebuilt-document": "document",
    "prebuilt-image": "image",
    "prebuilt-audio": "audio",
    "prebuilt-video": "video",
}

# Cache of known prebuilt analyzer name → modality (avoids API call)
_KNOWN_PREBUILT_MODALITY: Dict[str, str] = {
    # Document-based prebuilts
    "prebuilt-documentSearch": "document",
    "prebuilt-layout": "document",
    "prebuilt-read": "document",
    "prebuilt-document": "document",
    "prebuilt-invoice": "document",
    "prebuilt-receipt": "document",
    "prebuilt-receipt.generic": "document",
    "prebuilt-receipt.hotel": "document",
    "prebuilt-idDocument": "document",
    "prebuilt-idDocument.generic": "document",
    "prebuilt-idDocument.passport": "document",
    "prebuilt-healthInsuranceCard.us": "document",
    "prebuilt-contract": "document",
    "prebuilt-creditCard": "document",
    "prebuilt-creditMemo": "document",
    "prebuilt-bankStatement.us": "document",
    "prebuilt-check.us": "document",
    "prebuilt-purchaseOrder": "document",
    "prebuilt-procurement": "document",
    "prebuilt-payStub.us": "document",
    "prebuilt-utilityBill": "document",
    "prebuilt-marriageCertificate.us": "document",
    "prebuilt-documentFieldSchema": "document",
    "prebuilt-documentFields": "document",
    # Tax prebuilts (all document-based)
    "prebuilt-tax.us": "document",
    "prebuilt-tax.us.w2": "document",
    "prebuilt-tax.us.w4": "document",
    "prebuilt-tax.us.1040": "document",
    # Mortgage prebuilts
    "prebuilt-mortgage.us": "document",
    "prebuilt-mortgage.us.1003": "document",
    "prebuilt-mortgage.us.closingDisclosure": "document",
    # Image-based prebuilts
    "prebuilt-image": "image",
    "prebuilt-imageSearch": "image",
    # Audio-based prebuilts
    "prebuilt-audio": "audio",
    "prebuilt-audioSearch": "audio",
    "prebuilt-callCenter": "audio",
    # Video-based prebuilts
    "prebuilt-video": "video",
    "prebuilt-videoSearch": "video",
    "prebuilt-videoSynopsis": "video",
}

def _resolve_analyzer_modality(client: Any, analyzer_id: str) -> str:
    """Resolve analyzer modality from cache or via get_analyzer() fallback.

    For known prebuilt-* names, returns the modality from
    ``_KNOWN_PREBUILT_MODALITY`` without an API call.  For unknown
    prebuilt-* names or custom analyzers, calls ``get_analyzer()``
    to inspect ``base_analyzer_id``.

    Args:
        client: A ``ContentUnderstandingClient`` instance.
        analyzer_id: The analyzer ID to resolve.

    Returns:
        Modality string ("document", "image", "audio", or "video").

    Raises:
        ValueError: If ``get_analyzer()`` fails.
    """
    # Known prebuilt — use cache, no API call
    if analyzer_id in _KNOWN_PREBUILT_MODALITY:
        return _KNOWN_PREBUILT_MODALITY[analyzer_id]

    # Unknown prebuilt or custom analyzer — call get_analyzer()
    try:
        analyzer_info = client.get_analyzer(analyzer_id)
    except Exception as exc:
        raise ValueError(f"Failed to resolve analyzer '{analyzer_id}': {exc}") from exc

    if analyzer_info.base_analyzer_id:
        return _BASE_TO_MODALITY.get(analyzer_info.base_analyzer_id, "document")
    return "document"

def _is_analyzer_compatible(file_modality: str, analyzer_modality: str) -> bool:
    """Return True when an analyzer modality can process a file modality."""
    if analyzer_modality == "document":
        return file_modality in {"document", "image"}
    return file_modality == analyzer_modality

# ---------------------------------------------------------------------------
# Converter
# ---------------------------------------------------------------------------

class ContentUnderstandingConverter(DocumentConverter):
    """Converts files using Azure Content Understanding.

    Provides high-quality document, image, audio, and video conversion
    with structured field extraction via YAML front matter.
    """

    def __init__(
        self,
        *,
        endpoint: str,
        credential: AzureKeyCredential | TokenCredential | None = None,
        analyzer_id: Optional[str] = None,
        file_types: Optional[List[ContentUnderstandingFileType]] = None,
    ):
        """Initialize the Content Understanding converter.

        Args:
            endpoint: CU resource endpoint URL.
            credential: Explicit credential. If None, falls back to
                AZURE_API_KEY env var, then DefaultAzureCredential.
            analyzer_id: Custom analyzer for compatible file types.
                When set, the converter checks the analyzer's base modality
                (via get_analyzer() at init) and routes only compatible
                file types to it. Incompatible modalities auto-route to
                default prebuilts. If None, auto-selects by extension/MIME.
            file_types: Which file types to handle. If None, uses the
                default set (all supported formats).
        """
        super().__init__()

        # Raise if dependencies are missing
        if _dependency_exc_info is not None:
            raise MissingDependencyException(
                "ContentUnderstandingConverter requires the optional dependency "
                "[az-content-understanding] (or [all]) to be installed. "
                "E.g., `pip install 'markitdown[az-content-understanding]'`"
            ) from _dependency_exc_info[
                1
            ].with_traceback(  # type: ignore[union-attr]
                _dependency_exc_info[2]
            )

        self._file_types = file_types if file_types is not None else _ALL_FILE_TYPES
        self._analyzer_id = analyzer_id
        self._analyzer_modality: Optional[str] = None

        # Resolve credential
        if credential is None:
            api_key = os.environ.get("AZURE_API_KEY")
            if api_key is not None:
                credential = AzureKeyCredential(api_key)
            else:
                credential = DefaultAzureCredential()

        # User agent for telemetry
        try:
            from ..__about__ import __version__
        except ImportError:
            __version__ = "unknown"
        user_agent = f"markitdown-cu/{__version__}"

        # Create CU client
        self._client = ContentUnderstandingClient(
            endpoint=endpoint,
            credential=credential,
            user_agent_policy=UserAgentPolicy(user_agent=user_agent),
        )

        # Smart routing: resolve analyzer modality at init (at most one API call)
        if self._analyzer_id is not None:
            self._analyzer_modality = _resolve_analyzer_modality(
                self._client, self._analyzer_id
            )

    def accepts(
        self,
        file_stream: BinaryIO,
        stream_info: StreamInfo,
        **kwargs: Any,
    ) -> bool:
        """Return True if the file type is in the configured set."""
        return _detect_file_type(stream_info, self._file_types) is not None

    def convert(
        self,
        file_stream: BinaryIO,
        stream_info: StreamInfo,
        **kwargs: Any,
    ) -> DocumentConverterResult:
        """Convert the file using CU and return Markdown with YAML front matter."""

        # 1. Determine analyzer_id (smart routing: check modality)
        file_type = _detect_file_type(stream_info, self._file_types)
        if file_type is None:
            raise ValueError(
                "Unsupported file type for Content Understanding conversion."
            )
        file_modality = _get_modality(file_type)

        if (
            self._analyzer_id is not None
            and self._analyzer_modality is not None
            and _is_analyzer_compatible(file_modality, self._analyzer_modality)
        ):
            analyzer_id = self._analyzer_id
        else:
            analyzer_id = _PREBUILT_ANALYZERS.get(
                file_modality, "prebuilt-documentSearch"
            )

        # 2. Read file bytes and determine MIME type
        file_bytes = file_stream.read()
        content_type = _content_type_for(file_type, stream_info.mimetype)

        # 3. Call CU SDK
        poller = self._client.begin_analyze_binary(
            analyzer_id=analyzer_id,
            binary_input=file_bytes,
            content_type=content_type,
        )

        # 4. Block on result
        result = poller.result()

        # 5. Format output using to_llm_input()
        text = to_llm_input(result)

        # 6. Return
        return DocumentConverterResult(markdown=text)


package io.unitycatalog.server.sdk.tempcredential;

import static org.assertj.core.api.Assertions.assertThatThrownBy;

import io.unitycatalog.client.ApiException;
import io.unitycatalog.client.api.TemporaryCredentialsApi;
import io.unitycatalog.client.model.GenerateTemporaryPathCredential;
import io.unitycatalog.client.model.PathOperation;
import io.unitycatalog.client.model.TemporaryCredentials;
import io.unitycatalog.server.base.BaseCRUDTestWithMockCredentials;
import io.unitycatalog.server.base.ServerConfig;
import io.unitycatalog.server.base.catalog.CatalogOperations;
import io.unitycatalog.server.base.schema.SchemaOperations;
import io.unitycatalog.server.exception.ErrorCode;
import io.unitycatalog.server.sdk.catalog.SdkCatalogOperations;
import io.unitycatalog.server.sdk.schema.SdkSchemaOperations;
import io.unitycatalog.server.utils.TestUtils;
import java.util.List;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;

public class SdkTemporaryPathCredentialTest extends BaseCRUDTestWithMockCredentials {
  private TemporaryCredentialsApi temporaryCredentialsApi;

  @Override
  protected CatalogOperations createCatalogOperations(ServerConfig serverConfig) {
    return new SdkCatalogOperations(TestUtils.createApiClient(serverConfig));
  }

  @Override
  protected SchemaOperations createSchemaOperations(ServerConfig serverConfig) {
    return new SdkSchemaOperations(TestUtils.createApiClient(serverConfig));
  }

  @BeforeEach
  @Override
  public void setUp() {
    super.setUp();
    temporaryCredentialsApi = new TemporaryCredentialsApi(TestUtils.createApiClient(serverConfig));
  }

  @ParameterizedTest
  @MethodSource("getArgumentsForParameterizedTests")
  public void testGenerateTemporaryCredentialsWhereConfIsProvided(
      String scheme, boolean isConfiguredPath) throws ApiException {
    String url = getTestCloudPath(scheme, isConfiguredPath);
    GenerateTemporaryPathCredential generateTemporaryPathCredential =
        new GenerateTemporaryPathCredential().url(url).operation(PathOperation.PATH_READ);
    if (isConfiguredPath) {
      TemporaryCredentials temporaryCredentials =
          temporaryCredentialsApi.generateTemporaryPathCredentials(generateTemporaryPathCredential);
      assertTemporaryCredentials(temporaryCredentials, scheme, url);
    } else {
      assertThatThrownBy(
              () ->
                  temporaryCredentialsApi.generateTemporaryPathCredentials(
                      generateTemporaryPathCredential))
          .isInstanceOf(ApiException.class);
    }
  }

  @Test
  public void testGenerateAwsTemporaryCredentialsFromMasterRole() throws ApiException {
    for (String url : List.of(AWS_EXTERNAL_LOCATION_PATH, AWS_EXTERNAL_LOCATION_PATH + "/table1")) {
      GenerateTemporaryPathCredential generateTemporaryPathCredential =
          new GenerateTemporaryPathCredential().url(url).operation(PathOperation.PATH_READ_WRITE);
      TemporaryCredentials temporaryCredentials =
          temporaryCredentialsApi.generateTemporaryPathCredentials(generateTemporaryPathCredential);
      EchoAwsStsClient.assertAwsCredential(temporaryCredentials);
    }
    // Should fail because the path is not covered by external location
    TestUtils.assertApiException(
        () ->
            temporaryCredentialsApi.generateTemporaryPathCredentials(
                new GenerateTemporaryPathCredential()
                    .url(AWS_EXTERNAL_LOCATION_PARENT_PATH)
                    .operation(PathOperation.PATH_READ_WRITE)),
        ErrorCode.FAILED_PRECONDITION,
        "S3 bucket configuration not found");
  }
}


"""
Convenient access to LLM models.

Usage:
    from browser_use import llm

    # Simple model access
    model = llm.azure_gpt_4_1_mini
    model = llm.openai_gpt_4o
    model = llm.google_gemini_2_5_pro
    model = llm.bu_latest  # or bu_1_0, bu_2_0
"""

import os
from typing import TYPE_CHECKING

from browser_use.llm.azure.chat import ChatAzureOpenAI
from browser_use.llm.browser_use.chat import ChatBrowserUse
from browser_use.llm.cerebras.chat import ChatCerebras
from browser_use.llm.google.chat import ChatGoogle
from browser_use.llm.mistral.chat import ChatMistral
from browser_use.llm.openai.chat import ChatOpenAI

# Optional OCI import
try:
	from browser_use.llm.oci_raw.chat import ChatOCIRaw

	OCI_AVAILABLE = True
except ImportError:
	ChatOCIRaw = None
	OCI_AVAILABLE = False

if TYPE_CHECKING:
	from browser_use.llm.base import BaseChatModel

# Type stubs for IDE autocomplete
openai_gpt_4o: 'BaseChatModel'
openai_gpt_4o_mini: 'BaseChatModel'
openai_gpt_4_1_mini: 'BaseChatModel'
openai_o1: 'BaseChatModel'
openai_o1_mini: 'BaseChatModel'
openai_o1_pro: 'BaseChatModel'
openai_o3: 'BaseChatModel'
openai_o3_mini: 'BaseChatModel'
openai_o3_pro: 'BaseChatModel'
openai_o4_mini: 'BaseChatModel'
openai_gpt_5: 'BaseChatModel'
openai_gpt_5_mini: 'BaseChatModel'
openai_gpt_5_nano: 'BaseChatModel'

azure_gpt_4o: 'BaseChatModel'
azure_gpt_4o_mini: 'BaseChatModel'
azure_gpt_4_1_mini: 'BaseChatModel'
azure_o1: 'BaseChatModel'
azure_o1_mini: 'BaseChatModel'
azure_o1_pro: 'BaseChatModel'
azure_o3: 'BaseChatModel'
azure_o3_mini: 'BaseChatModel'
azure_o3_pro: 'BaseChatModel'
azure_gpt_5: 'BaseChatModel'
azure_gpt_5_mini: 'BaseChatModel'

google_gemini_2_0_flash: 'BaseChatModel'
google_gemini_2_0_pro: 'BaseChatModel'
google_gemini_2_5_pro: 'BaseChatModel'
google_gemini_2_5_flash: 'BaseChatModel'
google_gemini_2_5_flash_lite: 'BaseChatModel'
mistral_large: 'BaseChatModel'
mistral_medium: 'BaseChatModel'
mistral_small: 'BaseChatModel'
codestral: 'BaseChatModel'
pixtral_large: 'BaseChatModel'

anthropic_claude_sonnet_4_0: 'BaseChatModel'
anthropic_claude_fable_5: 'BaseChatModel'
anthropic_claude_3_5_sonnet_latest: 'BaseChatModel'
anthropic_claude_3_5_haiku_latest: 'BaseChatModel'

cerebras_llama3_1_8b: 'BaseChatModel'
cerebras_llama3_3_70b: 'BaseChatModel'
cerebras_gpt_oss_120b: 'BaseChatModel'
cerebras_llama_4_scout_17b_16e_instruct: 'BaseChatModel'
cerebras_llama_4_maverick_17b_128e_instruct: 'BaseChatModel'
cerebras_qwen_3_32b: 'BaseChatModel'
cerebras_qwen_3_235b_a22b_instruct_2507: 'BaseChatModel'
cerebras_qwen_3_235b_a22b_thinking_2507: 'BaseChatModel'
cerebras_qwen_3_coder_480b: 'BaseChatModel'

bu_latest: 'BaseChatModel'
bu_1_0: 'BaseChatModel'
bu_2_0: 'BaseChatModel'

def get_llm_by_name(model_name: str):
	"""
	Factory function to create LLM instances from string names with API keys from environment.

	Args:
	    model_name: String name like 'azure_gpt_4_1_mini', 'openai_gpt_4o', etc.

	Returns:
	    LLM instance with API keys from environment variables

	Raises:
	    ValueError: If model_name is not recognized
	"""
	if not model_name:
		raise ValueError('Model name cannot be empty')

	# Handle top-level Mistral aliases without provider prefix
	mistral_aliases = {
		'mistral_large': 'mistral-large-latest',
		'mistral_medium': 'mistral-medium-latest',
		'mistral_small': 'mistral-small-latest',
		'codestral': 'codestral-latest',
		'pixtral_large': 'pixtral-large-latest',
	}
	if model_name in mistral_aliases:
		api_key = os.getenv('MISTRAL_API_KEY')
		base_url = os.getenv('MISTRAL_BASE_URL', 'https://api.mistral.ai/v1')
		return ChatMistral(model=mistral_aliases[model_name], api_key=api_key, base_url=base_url)

	# Parse model name
	parts = model_name.split('_', 1)
	if len(parts) < 2:
		raise ValueError(f"Invalid model name format: '{model_name}'. Expected format: 'provider_model_name'")

	provider = parts[0]
	model_part = parts[1]

	# Convert underscores back to dots/dashes for actual model names
	if 'gpt_4_1_mini' in model_part:
		model = model_part.replace('gpt_4_1_mini', 'gpt-4.1-mini')
	elif 'gpt_4o_mini' in model_part:
		model = model_part.replace('gpt_4o_mini', 'gpt-4o-mini')
	elif 'gpt_4o' in model_part:
		model = model_part.replace('gpt_4o', 'gpt-4o')
	elif 'gemini_2_0' in model_part:
		model = model_part.replace('gemini_2_0', 'gemini-2.0').replace('_', '-')
	elif 'gemini_2_5' in model_part:
		model = model_part.replace('gemini_2_5', 'gemini-2.5').replace('_', '-')
	elif 'llama3_1' in model_part:
		model = model_part.replace('llama3_1', 'llama3.1').replace('_', '-')
	elif 'llama3_3' in model_part:
		model = model_part.replace('llama3_3', 'llama-3.3').replace('_', '-')
	elif 'llama_4_scout' in model_part:
		model = model_part.replace('llama_4_scout', 'llama-4-scout').replace('_', '-')
	elif 'llama_4_maverick' in model_part:
		model = model_part.replace('llama_4_maverick', 'llama-4-maverick').replace('_', '-')
	elif 'gpt_oss_120b' in model_part:
		model = model_part.replace('gpt_oss_120b', 'gpt-oss-120b')
	elif 'qwen_3_32b' in model_part:
		model = model_part.replace('qwen_3_32b', 'qwen-3-32b')
	elif 'qwen_3_235b_a22b_instruct' in model_part:
		if model_part.endswith('_2507'):
			model = model_part.replace('qwen_3_235b_a22b_instruct_2507', 'qwen-3-235b-a22b-instruct-2507')
		else:
			model = model_part.replace('qwen_3_235b_a22b_instruct', 'qwen-3-235b-a22b-instruct-2507')
	elif 'qwen_3_235b_a22b_thinking' in model_part:
		if model_part.endswith('_2507'):
			model = model_part.replace('qwen_3_235b_a22b_thinking_2507', 'qwen-3-235b-a22b-thinking-2507')
		else:
			model = model_part.replace('qwen_3_235b_a22b_thinking', 'qwen-3-235b-a22b-thinking-2507')
	elif 'qwen_3_coder_480b' in model_part:
		model = model_part.replace('qwen_3_coder_480b', 'qwen-3-coder-480b')
	else:
		model = model_part.replace('_', '-')

	# OpenAI Models
	if provider == 'openai':
		api_key = os.getenv('OPENAI_API_KEY')
		return ChatOpenAI(model=model, api_key=api_key)

	# Azure OpenAI Models
	elif provider == 'azure':
		api_key = os.getenv('AZURE_OPENAI_KEY') or os.getenv('AZURE_OPENAI_API_KEY')
		azure_endpoint = os.getenv('AZURE_OPENAI_ENDPOINT')
		return ChatAzureOpenAI(model=model, api_key=api_key, azure_endpoint=azure_endpoint)

	# Google Models
	elif provider == 'google':
		api_key = os.getenv('GOOGLE_API_KEY')
		return ChatGoogle(model=model, api_key=api_key)

	# Anthropic Models
	elif provider == 'anthropic':
		from browser_use.llm.anthropic.chat import ChatAnthropic

		api_key = os.getenv('ANTHROPIC_API_KEY')
		return ChatAnthropic(model=model, api_key=api_key)

	# Mistral Models
	elif provider == 'mistral':
		api_key = os.getenv('MISTRAL_API_KEY')
		base_url = os.getenv('MISTRAL_BASE_URL', 'https://api.mistral.ai/v1')
		mistral_map = {
			'large': 'mistral-large-latest',
			'medium': 'mistral-medium-latest',
			'small': 'mistral-small-latest',
			'codestral': 'codestral-latest',
			'pixtral-large': 'pixtral-large-latest',
		}
		normalized_model_part = model_part.replace('_', '-')
		resolved_model = mistral_map.get(normalized_model_part, model.replace('_', '-'))
		return ChatMistral(model=resolved_model, api_key=api_key, base_url=base_url)

	# OCI Models
	elif provider == 'oci':
		# OCI requires more complex configuration that can't be easily inferred from env vars
		# Users should use ChatOCIRaw directly with proper configuration
		raise ValueError('OCI models require manual configuration. Use ChatOCIRaw directly with your OCI credentials.')

	# Cerebras Models
	elif provider == 'cerebras':
		api_key = os.getenv('CEREBRAS_API_KEY')
		return ChatCerebras(model=model, api_key=api_key)

	# Browser Use Models
	elif provider == 'bu':
		# Handle bu_latest -> bu-latest conversion (need to prepend 'bu-' back)
		model = f'bu-{model_part.replace("_", "-")}'
		api_key = os.getenv('BROWSER_USE_API_KEY')
		return ChatBrowserUse(model=model, api_key=api_key)

	else:
		available_providers = ['openai', 'azure', 'google', 'anthropic', 'mistral', 'oci', 'cerebras', 'bu']
		raise ValueError(f"Unknown provider: '{provider}'. Available providers: {', '.join(available_providers)}")

# Pre-configured model instances (lazy loaded via __getattr__)
def __getattr__(name: str) -> 'BaseChatModel':
	"""Create model instances on demand with API keys from environment."""
	# Handle chat classes first
	if name == 'ChatOpenAI':
		return ChatOpenAI  # type: ignore
	elif name == 'ChatAzureOpenAI':
		return ChatAzureOpenAI  # type: ignore
	elif name == 'ChatGoogle':
		return ChatGoogle  # type: ignore

	elif name == 'ChatMistral':
		return ChatMistral  # type: ignore

	elif name == 'ChatOCIRaw':
		if not OCI_AVAILABLE:
			raise ImportError('OCI integration not available. Install with: pip install "browser-use[oci]"')
		return ChatOCIRaw  # type: ignore
	elif name == 'ChatCerebras':
		return ChatCerebras  # type: ignore
	elif name == 'ChatBrowserUse':
		return ChatBrowserUse  # type: ignore

	# Handle model instances - these are the main use case
	try:
		return get_llm_by_name(name)
	except ValueError:
		raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

# Export all classes and preconfigured instances, conditionally including ChatOCIRaw
__all__ = [
	'ChatOpenAI',
	'ChatAzureOpenAI',
	'ChatGoogle',
	'ChatMistral',
	'ChatCerebras',
	'ChatBrowserUse',
]

if OCI_AVAILABLE:
	__all__.append('ChatOCIRaw')

__all__ += [
	'get_llm_by_name',
	# OpenAI instances - created on demand
	'openai_gpt_4o',
	'openai_gpt_4o_mini',
	'openai_gpt_4_1_mini',
	'openai_o1',
	'openai_o1_mini',
	'openai_o1_pro',
	'openai_o3',
	'openai_o3_mini',
	'openai_o3_pro',
	'openai_o4_mini',
	'openai_gpt_5',
	'openai_gpt_5_mini',
	'openai_gpt_5_nano',
	# Azure instances - created on demand
	'azure_gpt_4o',
	'azure_gpt_4o_mini',
	'azure_gpt_4_1_mini',
	'azure_o1',
	'azure_o1_mini',
	'azure_o1_pro',
	'azure_o3',
	'azure_o3_mini',
	'azure_o3_pro',
	'azure_gpt_5',
	'azure_gpt_5_mini',
	# Google instances - created on demand
	'google_gemini_2_0_flash',
	'google_gemini_2_0_pro',
	'google_gemini_2_5_pro',
	'google_gemini_2_5_flash',
	'google_gemini_2_5_flash_lite',
	# Anthropic instances - created on demand
	'anthropic_claude_sonnet_4_0',
	'anthropic_claude_fable_5',
	'anthropic_claude_3_5_sonnet_latest',
	'anthropic_claude_3_5_haiku_latest',
	# Mistral instances - created on demand
	'mistral_large',
	'mistral_medium',
	'mistral_small',
	'codestral',
	'pixtral_large',
	# Cerebras instances - created on demand
	'cerebras_llama3_1_8b',
	'cerebras_llama3_3_70b',
	'cerebras_gpt_oss_120b',
	'cerebras_llama_4_scout_17b_16e_instruct',
	'cerebras_llama_4_maverick_17b_128e_instruct',
	'cerebras_qwen_3_32b',
	'cerebras_qwen_3_235b_a22b_instruct_2507',
	'cerebras_qwen_3_235b_a22b_thinking_2507',
	'cerebras_qwen_3_coder_480b',
	# Browser Use instances - created on demand
	'bu_latest',
	'bu_1_0',
	'bu_2_0',
]

# NOTE: OCI backend is optional. The try/except ImportError and conditional __all__ are required
# so this module can be imported without browser-use[oci] installed.


package io.unitycatalog.server.sdk.access;

import static io.unitycatalog.server.utils.TestUtils.assertPermissionDenied;
import static org.assertj.core.api.Assertions.assertThat;

import io.unitycatalog.client.api.ModelVersionsApi;
import io.unitycatalog.client.api.RegisteredModelsApi;
import io.unitycatalog.client.model.CreateModelVersion;
import io.unitycatalog.client.model.CreateRegisteredModel;
import io.unitycatalog.client.model.ModelVersionInfo;
import io.unitycatalog.client.model.RegisteredModelInfo;
import io.unitycatalog.client.model.SecurableType;
import io.unitycatalog.client.model.UpdateModelVersion;
import io.unitycatalog.client.model.UpdateRegisteredModel;
import io.unitycatalog.server.base.ServerConfig;
import io.unitycatalog.server.persist.model.Privileges;
import io.unitycatalog.server.utils.TestUtils;
import java.util.List;
import lombok.SneakyThrows;
import org.junit.jupiter.api.Test;

/**
 * SDK-based access control tests for Registered Model and Model Version CRUD operations.
 *
 *
This test class verifies:
 *
 *

 *
Model creation requires CREATE MODEL permission on schema
 *
Model get requires ownership
 *
Model list is filtered based on ownership
 *
Model update requires ownership
 *
Model version creation requires ownership of parent model
 *
Model version get requires ownership
 *
Model version list requires ownership
 *
Model version update requires ownership
 *
Model version delete requires ownership
 *
Model delete requires ownership
 *

 */
public class SdkModelAccessControlCRUDTest extends SdkAccessControlBaseCRUDTest {

  @Test
  @SneakyThrows
  public void testModelAccess() {
    createCommonTestUsers();
    setupCommonCatalogAndSchema();

    // Create API clients for different users
    ServerConfig principal1Config = createTestUserServerConfig(PRINCIPAL_1);
    ServerConfig principal2Config = createTestUserServerConfig(PRINCIPAL_2);

    RegisteredModelsApi principal1ModelsApi =
        new RegisteredModelsApi(TestUtils.createApiClient(principal1Config));
    RegisteredModelsApi principal2ModelsApi =
        new RegisteredModelsApi(TestUtils.createApiClient(principal2Config));
    ModelVersionsApi principal1VersionsApi =
        new ModelVersionsApi(TestUtils.createApiClient(principal1Config));
    ModelVersionsApi principal2VersionsApi =
        new ModelVersionsApi(TestUtils.createApiClient(principal2Config));

    // Grant USE SCHEMA and CREATE MODEL to principal-1
    grantPermissions(PRINCIPAL_1, SecurableType.SCHEMA, "cat_pr1.sch_pr1", Privileges.USE_SCHEMA);
    grantPermissions(PRINCIPAL_1, SecurableType.SCHEMA, "cat_pr1.sch_pr1", Privileges.CREATE_MODEL);

    // Grant USE SCHEMA to principal-2
    grantPermissions(PRINCIPAL_2, SecurableType.SCHEMA, "cat_pr1.sch_pr1", Privileges.USE_SCHEMA);

    // TEST: Create registered model as principal-1 - should succeed
    CreateRegisteredModel createModel =
        new CreateRegisteredModel().name("mod_pr1").catalogName("cat_pr1").schemaName("sch_pr1");
    RegisteredModelInfo modelInfo = principal1ModelsApi.createRegisteredModel(createModel);
    assertThat(modelInfo).isNotNull();
    assertThat(modelInfo.getName()).isEqualTo("mod_pr1");

    // TEST: Get registered model as principal-1 (owner) - should succeed
    RegisteredModelInfo getModelInfo =
        principal1ModelsApi.getRegisteredModel("cat_pr1.sch_pr1.mod_pr1");
    assertThat(getModelInfo).isNotNull();

    // TEST: Get registered model as principal-2 (not owner) - should fail
    assertPermissionDenied(() -> principal2ModelsApi.getRegisteredModel("cat_pr1.sch_pr1.mod_pr1"));

    // TEST: List registered models as principal-1 - should see owned model
    List<RegisteredModelInfo> principal1Models =
        listAllRegisteredModels(principal1ModelsApi, "cat_pr1", "sch_pr1");
    assertThat(principal1Models).hasSize(1);

    // TEST: List registered models as principal-2 - should see empty list
    List<RegisteredModelInfo> principal2Models =
        listAllRegisteredModels(principal2ModelsApi, "cat_pr1", "sch_pr1");
    assertThat(principal2Models).isEmpty();

    // TEST: Update registered model as principal-1 (owner) - should succeed
    UpdateRegisteredModel updateModel = new UpdateRegisteredModel().comment("hello");
    RegisteredModelInfo updatedModel =
        principal1ModelsApi.updateRegisteredModel("cat_pr1.sch_pr1.mod_pr1", updateModel);
    assertThat(updatedModel.getComment()).isEqualTo("hello");

    // TEST: Update registered model as principal-2 (not owner) - should fail
    UpdateRegisteredModel updateModel2 = new UpdateRegisteredModel().comment("hello2");
    assertPermissionDenied(
        () -> principal2ModelsApi.updateRegisteredModel("cat_pr1.sch_pr1.mod_pr1", updateModel2));

    // TEST: Create model version as principal-1 (owner) - should succeed
    CreateModelVersion createVersion =
        new CreateModelVersion()
            .catalogName("cat_pr1")
            .schemaName("sch_pr1")
            .modelName("mod_pr1")
            .source("model_source");
    ModelVersionInfo versionInfo = principal1VersionsApi.createModelVersion(createVersion);
    assertThat(versionInfo).isNotNull();
    assertThat(versionInfo.getVersion()).isEqualTo(1L);

    // TEST: Create model version as principal-2 (not owner) - should fail
    CreateModelVersion createVersion2 =
        new CreateModelVersion()
            .catalogName("cat_pr1")
            .schemaName("sch_pr1")
            .modelName("mod_pr1")
            .source("model_source");
    assertPermissionDenied(() -> principal2VersionsApi.createModelVersion(createVersion2));

    // TEST: Get model version as principal-1 (owner) - should succeed
    ModelVersionInfo getVersionInfo =
        principal1VersionsApi.getModelVersion("cat_pr1.sch_pr1.mod_pr1", 1L);
    assertThat(getVersionInfo).isNotNull();

    // TEST: Get model version as principal-2 (not owner) - should fail
    assertPermissionDenied(
        () -> principal2VersionsApi.getModelVersion("cat_pr1.sch_pr1.mod_pr1", 1L));

    // TEST: List model versions as principal-1 - should see all versions
    List<ModelVersionInfo> principal1Versions =
        listAllModelVersions(principal1VersionsApi, "cat_pr1.sch_pr1.mod_pr1");
    assertThat(principal1Versions).hasSize(1);

    // TEST: List model versions as principal-2 - should fail
    assertPermissionDenied(
        () -> listAllModelVersions(principal2VersionsApi, "cat_pr1.sch_pr1.mod_pr1"));

    // TEST: Update model version as principal-1 (owner) - should succeed
    UpdateModelVersion updateVersion = new UpdateModelVersion().comment("hello");
    ModelVersionInfo updatedVersion =
        principal1VersionsApi.updateModelVersion("cat_pr1.sch_pr1.mod_pr1", 1L, updateVersion);
    assertThat(updatedVersion.getComment()).isEqualTo("hello");

    // TEST: Update model version as principal-2 (not owner) - should fail
    UpdateModelVersion updateVersion2 = new UpdateModelVersion().comment("hello2");
    assertPermissionDenied(
        () ->
            principal2VersionsApi.updateModelVersion(
                "cat_pr1.sch_pr1.mod_pr1", 1L, updateVersion2));

    // TEST: Delete model version as principal-2 (not owner) - should fail
    assertPermissionDenied(
        () -> principal2VersionsApi.deleteModelVersion("cat_pr1.sch_pr1.mod_pr1", 1L));

    // TEST: Delete model version as principal-1 (owner) - should succeed
    principal1VersionsApi.deleteModelVersion("cat_pr1.sch_pr1.mod_pr1", 1L);

    // Verify deletion
    List<ModelVersionInfo> versionsAfterDelete =
        listAllModelVersions(principal1VersionsApi, "cat_pr1.sch_pr1.mod_pr1");
    assertThat(versionsAfterDelete).isEmpty();

    // TEST: Delete registered model as principal-2 (not owner) - should fail
    assertPermissionDenied(
        () -> principal2ModelsApi.deleteRegisteredModel("cat_pr1.sch_pr1.mod_pr1", false));

    // TEST: Delete registered model as principal-1 (owner) - should succeed
    principal1ModelsApi.deleteRegisteredModel("cat_pr1.sch_pr1.mod_pr1", false);

    // Verify deletion
    List<RegisteredModelInfo> modelsAfterDelete =
        listAllRegisteredModels(principal1ModelsApi, "cat_pr1", "sch_pr1");
    assertThat(modelsAfterDelete).isEmpty();
  }
}


from typing import Optional, List, Any

MISSING_DEPENDENCY_MESSAGE = """{converter} recognized the input as a potential {extension} file, but the dependencies needed to read {extension} files have not been installed. To resolve this error, include the optional dependency [{feature}] or [all] when installing MarkItDown. For example:

* pip install 'markitdown[{feature}]'
* pip install 'markitdown[all]'
* pip install 'markitdown[{feature}, ...]'
* etc."""

class MarkItDownException(Exception):
    """
    Base exception class for MarkItDown.
    """

    pass

class MissingDependencyException(MarkItDownException):
    """
    Converters shipped with MarkItDown may depend on optional
    dependencies. This exception is thrown when a converter's
    convert() method is called, but the required dependency is not
    installed. This is not necessarily a fatal error, as the converter
    will simply be skipped (an error will bubble up only if no other
    suitable converter is found).

    Error messages should clearly indicate which dependency is missing.
    """

    pass

class UnsupportedFormatException(MarkItDownException):
    """
    Thrown when no suitable converter was found for the given file.
    """

    pass

class FailedConversionAttempt(object):
    """
    Represents a single attempt to convert a file.
    """

    def __init__(self, converter: Any, exc_info: Optional[tuple] = None):
        self.converter = converter
        self.exc_info = exc_info

class FileConversionException(MarkItDownException):
    """
    Thrown when a suitable converter was found, but the conversion
    process fails for any reason.
    """

    def __init__(
        self,
        message: Optional[str] = None,
        attempts: Optional[List[FailedConversionAttempt]] = None,
    ):
        self.attempts = attempts

        if message is None:
            if attempts is None:
                message = "File conversion failed."
            else:
                message = f"File conversion failed after {len(attempts)} attempts:\n"
                for attempt in attempts:
                    if attempt.exc_info is None:
                        message += f" -  {type(attempt.converter).__name__} provided no execution info."
                    else:
                        message += f" - {type(attempt.converter).__name__} threw {attempt.exc_info[0].__name__} with message: {attempt.exc_info[1]}\n"

        super().__init__(message)


"use client"

import * as React from "react"
import { Command as CommandPrimitive } from "cmdk"

import { cn } from "#/lib/utils.ts"
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
} from "#/components/ui/dialog.tsx"
import {
  InputGroup,
  InputGroupAddon,
} from "#/components/ui/input-group.tsx"
import { HugeiconsIcon } from "@hugeicons/react"
import { SearchIcon, Tick02Icon } from "@hugeicons/core-free-icons"

function Command({
  className,
  ...props
}: React.ComponentProps<typeof CommandPrimitive>) {
  return (
    <CommandPrimitive
      data-slot="command"
      className={cn(
        "flex size-full flex-col overflow-hidden rounded-xl bg-popover p-1 text-popover-foreground",
        className
      )}
      {...props}
    />
  )
}

function CommandDialog({
  title = "Command Palette",
  description = "Search for a command to run...",
  children,
  className,
  showCloseButton = false,
  ...props
}: Omit<React.ComponentProps<typeof Dialog>, "children"> & {
  title?: string
  description?: string
  className?: string
  showCloseButton?: boolean
  children: React.ReactNode
}) {
  return (
    <Dialog {...props}>
      <DialogHeader className="sr-only">
        <DialogTitle>{title}</DialogTitle>
        <DialogDescription>{description}</DialogDescription>
      </DialogHeader>
      <DialogContent
        className={cn(
          "top-1/3 translate-y-0 overflow-hidden rounded-xl! p-0",
          className
        )}
        showCloseButton={showCloseButton}
      >
        {children}
      </DialogContent>
    </Dialog>
  )
}

function CommandInput({
  className,
  ...props
}: React.ComponentProps<typeof CommandPrimitive.Input>) {
  return (
    <div data-slot="command-input-wrapper" className="p-1 pb-0">
      <InputGroup className="h-8! bg-input/20 dark:bg-input/30">
        <CommandPrimitive.Input
          data-slot="command-input"
          className={cn(
            "w-full text-xs/relaxed outline-hidden disabled:cursor-not-allowed disabled:opacity-50",
            className
          )}
          {...props}
        />
        <InputGroupAddon>
          <HugeiconsIcon icon={SearchIcon} strokeWidth={2} className="size-3.5 shrink-0 opacity-50" />
        </InputGroupAddon>
      </InputGroup>
    </div>
  )
}

function CommandList({
  className,
  ...props
}: React.ComponentProps<typeof CommandPrimitive.List>) {
  return (
    <CommandPrimitive.List
      data-slot="command-list"
      className={cn(
        "no-scrollbar max-h-72 scroll-py-1 overflow-x-hidden overflow-y-auto outline-none",
        className
      )}
      {...props}
    />
  )
}

function CommandEmpty({
  className,
  ...props
}: React.ComponentProps<typeof CommandPrimitive.Empty>) {
  return (
    <CommandPrimitive.Empty
      data-slot="command-empty"
      className={cn("py-6 text-center text-xs/relaxed", className)}
      {...props}
    />
  )
}

function CommandGroup({
  className,
  ...props
}: React.ComponentProps<typeof CommandPrimitive.Group>) {
  return (
    <CommandPrimitive.Group
      data-slot="command-group"
      className={cn(
        "overflow-hidden p-1 text-foreground **:[[cmdk-group-heading]]:px-2.5 **:[[cmdk-group-heading]]:py-1.5 **:[[cmdk-group-heading]]:text-xs **:[[cmdk-group-heading]]:font-medium **:[[cmdk-group-heading]]:text-muted-foreground",
        className
      )}
      {...props}
    />
  )
}

function CommandSeparator({
  className,
  ...props
}: React.ComponentProps<typeof CommandPrimitive.Separator>) {
  return (
    <CommandPrimitive.Separator
      data-slot="command-separator"
      className={cn("-mx-1 my-1 h-px bg-border/50", className)}
      {...props}
    />
  )
}

function CommandItem({
  className,
  children,
  ...props
}: React.ComponentProps<typeof CommandPrimitive.Item>) {
  return (
    <CommandPrimitive.Item
      data-slot="command-item"
      className={cn(
        "group/command-item relative flex min-h-7 cursor-default items-center gap-2 rounded-md px-2.5 py-1.5 text-xs/relaxed outline-hidden select-none in-data-[slot=dialog-content]:rounded-md data-[disabled=true]:pointer-events-none data-[disabled=true]:opacity-50 data-selected:bg-muted data-selected:text-foreground [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-3.5 data-selected:*:[svg]:text-foreground",
        className
      )}
      {...props}
    >
      {children}
      <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto opacity-0 group-has-data-[slot=command-shortcut]/command-item:hidden group-data-[checked=true]/command-item:opacity-100" />
    </CommandPrimitive.Item>
  )
}

function CommandShortcut({
  className,
  ...props
}: React.ComponentProps<"span">) {
  return (
    <span
      data-slot="command-shortcut"
      className={cn(
        "ml-auto text-[0.625rem] tracking-widest text-muted-foreground group-data-selected/command-item:text-foreground",
        className
      )}
      {...props}
    />
  )
}

export {
  Command,
  CommandDialog,
  CommandInput,
  CommandList,
  CommandEmpty,
  CommandGroup,
  CommandItem,
  CommandShortcut,
  CommandSeparator,
}


import { describe, test, expect } from "vitest";
import { generateText, stepCountIs, tool } from "ai";
import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock";
import { z } from "zod";
import {
  resolveLibraryId,
  queryDocs,
  Context7Agent,
  SYSTEM_PROMPT,
  AGENT_PROMPT,
  RESOLVE_LIBRARY_ID_DESCRIPTION,
} from "./index";

const bedrock = createAmazonBedrock({
  region: process.env.AWS_REGION,
  apiKey: process.env.AWS_BEARER_TOKEN_BEDROCK,
});

describe("@upstash/context7-tools-ai-sdk", () => {
  describe("Tool structure", () => {
    test("resolveLibraryId() should return a tool object with correct structure", () => {
      const tool = resolveLibraryId();

      expect(tool).toBeDefined();
      expect(tool).toHaveProperty("execute");
      expect(tool).toHaveProperty("inputSchema");
      expect(tool).toHaveProperty("description");
      expect(tool.description).toContain("library");
    });

    test("queryDocs() should return a tool object with correct structure", () => {
      const tool = queryDocs();

      expect(tool).toBeDefined();
      expect(tool).toHaveProperty("execute");
      expect(tool).toHaveProperty("inputSchema");
      expect(tool).toHaveProperty("description");
      expect(tool.description).toContain("documentation");
    });

    test("tools should accept custom config", () => {
      const resolveTool = resolveLibraryId({
        apiKey: "ctx7sk-test-key",
      });

      const docsTool = queryDocs({
        apiKey: "ctx7sk-test-key",
      });

      expect(resolveTool).toHaveProperty("execute");
      expect(docsTool).toHaveProperty("execute");
    });
  });

  describe("Tool usage with generateText", () => {
    test("resolveLibraryId tool should be called when searching for a library", async () => {
      const result = await generateText({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        tools: {
          resolveLibraryId: resolveLibraryId(),
        },
        toolChoice: { type: "tool", toolName: "resolveLibraryId" },
        stopWhen: stepCountIs(2),
        prompt: "Search for 'react' library",
      });

      expect(result.toolCalls.length).toBeGreaterThan(0);
      expect(result.toolCalls[0].toolName).toBe("resolveLibraryId");
      expect(result.toolResults.length).toBeGreaterThan(0);
      const toolResult = result.toolResults[0] as unknown as { output: string };
      expect(typeof toolResult.output).toBe("string");
      expect(toolResult.output).toContain("Context7-compatible library ID");
    }, 30000);

    test("queryDocs tool should fetch documentation", async () => {
      const result = await generateText({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        tools: {
          queryDocs: queryDocs(),
        },
        toolChoice: { type: "tool", toolName: "queryDocs" },
        stopWhen: stepCountIs(2),
        prompt: "Fetch documentation for library ID '/facebook/react' about hooks",
      });

      expect(result.toolCalls.length).toBeGreaterThan(0);
      expect(result.toolCalls[0].toolName).toBe("queryDocs");
      expect(result.toolResults.length).toBeGreaterThan(0);
      const toolResult = result.toolResults[0] as unknown as { output: string };
      expect(typeof toolResult.output).toBe("string");
      expect(toolResult.output.length).toBeGreaterThan(0);
    }, 30000);

    test("both tools can work together in a multi-step flow", async () => {
      const result = await generateText({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        tools: {
          resolveLibraryId: resolveLibraryId(),
          queryDocs: queryDocs(),
        },
        stopWhen: stepCountIs(5),
        prompt:
          "First use resolveLibraryId to find the Next.js library, then use queryDocs to get documentation about routing",
      });

      const allToolCalls = result.steps.flatMap((step) => step.toolCalls);
      const toolNames = allToolCalls.map((call) => call.toolName);
      expect(toolNames).toContain("resolveLibraryId");
      expect(toolNames).toContain("queryDocs");
    }, 60000);
  });

  describe("Context7Agent class", () => {
    test("should create an agent instance with model", () => {
      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
      });

      expect(agent).toBeDefined();
      expect(agent).toHaveProperty("generate");
      expect(agent).toHaveProperty("stream");
    });

    test("should accept custom stopWhen condition", () => {
      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        stopWhen: stepCountIs(3),
      });

      expect(agent).toBeDefined();
    });

    test("should accept custom instructions", () => {
      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        instructions: "Custom instructions for testing",
      });

      expect(agent).toBeDefined();
    });

    test("should accept Context7 config options", () => {
      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        apiKey: "ctx7sk-test-key",
      });

      expect(agent).toBeDefined();
    });

    test("should accept additional tools alongside Context7 tools", () => {
      const customTool = tool({
        description: "A custom test tool",
        inputSchema: z.object({
          input: z.string().describe("Test input"),
        }),
        execute: async ({ input }) => ({ result: `processed: ${input}` }),
      });

      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        tools: {
          customTool,
        },
      });

      expect(agent).toBeDefined();
    });

    test("should generate response using agent workflow", async () => {
      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        stopWhen: stepCountIs(5),
      });

      const result = await agent.generate({
        prompt: "Find the React library and get documentation about hooks",
      });

      expect(result).toBeDefined();
      expect(result.steps.length).toBeGreaterThan(0);

      const allToolCalls = result.steps.flatMap((step) => step.toolCalls);
      const toolNames = allToolCalls.map((call) => call.toolName);
      expect(toolNames).toContain("resolveLibraryId");
    }, 60000);

    test("should include Context7 tools in generate result", async () => {
      const agent = new Context7Agent({
        model: bedrock("anthropic.claude-3-haiku-20240307-v1:0"),
        stopWhen: stepCountIs(5),
      });

      const result = await agent.generate({
        prompt:
          "Use resolveLibraryId to search for Next.js, then use queryDocs to get routing documentation",
      });

      expect(result).toBeDefined();

      const allToolCalls = result.steps.flatMap((step) => step.toolCalls);
      const toolNames = allToolCalls.map((call) => call.toolName);

      expect(toolNames).toContain("resolveLibraryId");
      expect(toolNames).toContain("queryDocs");
    }, 60000);
  });

  describe("Prompt exports", () => {
    test("should export SYSTEM_PROMPT", () => {
      expect(SYSTEM_PROMPT).toBeDefined();
      expect(typeof SYSTEM_PROMPT).toBe("string");
      expect(SYSTEM_PROMPT.length).toBeGreaterThan(0);
    });

    test("should export AGENT_PROMPT", () => {
      expect(AGENT_PROMPT).toBeDefined();
      expect(typeof AGENT_PROMPT).toBe("string");
      expect(AGENT_PROMPT).toContain("Context7");
    });

    test("should export RESOLVE_LIBRARY_ID_DESCRIPTION", () => {
      expect(RESOLVE_LIBRARY_ID_DESCRIPTION).toBeDefined();
      expect(typeof RESOLVE_LIBRARY_ID_DESCRIPTION).toBe("string");
      expect(RESOLVE_LIBRARY_ID_DESCRIPTION).toContain("library");
    });
  });
});