Python SDK — Code Examples

export COSMOS_ORG_ID="your-org-uuid"
export COSMOS_ACCESS_TOKEN="your-access-token"
export COSMOS_USER_SERVICE_URL="https://user.example.com"
export COSMOS_CONFIGURATOR_SERVICE_URL="https://configurator.example.com"
export COSMOS_ADK_SERVICE_URL="https://adk.example.com"
export COSMOS_STORAGE_SERVICE_URL="https://storage.example.com"

import os

from cosmos_sdk import CosmosClient
from cosmos_sdk import build_service_endpoints_from_env

client = CosmosClient.from_token(
    os.environ["COSMOS_ACCESS_TOKEN"],
    org_id=os.environ["COSMOS_ORG_ID"],
    services=build_service_endpoints_from_env(),
)

claims = client.auth.decode_token()
conversation_user_id = claims.sub or claims.user_id

from cosmos_sdk import ToolDefinition, ToolResult


async def search_orders(args: dict) -> ToolResult:
    customer_id = args["customer_id"]
    if not customer_id.startswith("C-"):
        # Sent to ADK as the tool_result error; the model sees this message
        return ToolResult(ok=False, error={"code": "TOOL_INVALID_INPUT", "message": "customer_id must start with C-"})
    return ToolResult(ok=True, data={"customer_id": customer_id, "orders": []})


client.tool_manager.register_tool(
    "search_orders",
    handler=search_orders,
    capabilities=["orders.search"],
    definition=ToolDefinition(
        name="search_orders",
        description="Search orders for a customer",
        input_schema={
            "type": "object",
            "properties": {"customer_id": {"type": "string"}},
            "required": ["customer_id"],
        },
    ),
)

from cosmos_sdk import ADKV2SendOptions, ADKV2TurnCallbacks

chunks: list[str] = []

result = await client.send_message(
    "Find the latest orders for customer C-100.",
    user_id=conversation_user_id,
    callbacks=ADKV2TurnCallbacks(
        on_chunk=lambda text: chunks.append(text),
        on_tool_call_args_delta=lambda name, path, delta: print(name, path, delta),
    ),
    options=ADKV2SendOptions(
        agent_id="agent-uuid",
        agent_version="1.0",
        session_state={"project_id": "project-uuid"},
        max_tool_roundtrips=20,
    ),
)

print("".join(chunks))
print(result.usage)

result = await client.send_message(
    "Draft a risk summary from the notes below.",
    user_id=conversation_user_id,
    options=ADKV2SendOptions(
        system_prompt="You are a concise risk analyst.",
        provider_id="azure_openai",
        model_id="gpt-4o",
        app_name="risk-reviewer",
    ),
)

agents = await client.search_agents({"query": "support"})
skills = await client.search_skills({"page_size": 50})
collections = await client.search_collections()

for skill in skills.items:
    print(skill["name"], skill.get("description", ""))

allowed = await client.authz.can({
    "actor": {"type": "user", "id": "user-uuid"},
    "action": "bridge.connect",
    "resource": {"type": "organization", "id": os.environ["COSMOS_ORG_ID"]},
    "context": {"channel": "jumpstart"},
})

if not allowed:
    raise PermissionError("Bridge access is not allowed for this user")

sources = await client.search_data_sources({"query": "warehouse", "engine": "sqlite"})

source = await client.create_data_source({
    "name": "Assessment SQLite",
    "engine": "sqlite",
    "connection_profile": {"path": "/dms/projects/acme/assessment.db"},
})

await client.sync_data_source_tool(source["id"])
print(await client.test_data_source(source["id"], query="SELECT 1"))
print(sources.total)

from cosmos_sdk import BuildInventoryOptions, build_inventory_prompt, build_resource_inventory

inventory = await build_resource_inventory(
    client,
    BuildInventoryOptions(
        skill_ids=["skill-uuid"],
        context_ids=["collection-uuid"],
    ),
)

prompt_section = build_inventory_prompt(
    kb_collections=inventory["kb_collections"],
    artefacts=inventory["artefacts"],
)

path = "/projects/acme/input/customer-notes.txt"
await client.storage.upload(path, b"Customer notes...", mime_type="text/plain")
item = await client.storage.resolve_path(f"{path}:latest")
if item:
    content = await client.storage.download_item(item["id"])
    print(content.decode("utf-8"))

result = await client.send_message(
    "Summarize the notes in the attached file.",
    user_id=conversation_user_id,
    options=ADKV2SendOptions(
        agent_id="agent-uuid",
        additional_instructions=f"Use the file at {path}:latest when needed.",
    ),
)

job = await client.storage.unzip_item("archive-item-uuid")
sync_result = await client.storage.unzip_item("archive-item-uuid", sync=True)
print(job, sync_result)

if client.ragflow:
    results = await client.ragflow.search(
        "collection-uuid",
        "What is the refund policy?",
        top_k=5,
        search_type="hybrid",
    )
    for result in results:
        print(result)

models = await client.fetch_org_models()
print([model for model in models if model.get("provider") == "azure_openai"])

# Register your own provider object that implements the LLMProvider protocol,
# then route calls through client.llm.complete().
client.llm.register_provider(my_provider)
response = await client.llm.complete(
    messages=[{"role": "user", "content": "Write one sentence."}],
)

from cosmos_sdk import CosmosApiError, CosmosAuthError

try:
    await client.get_agent("missing-agent")
except CosmosApiError as exc:
    if exc.status_code == 404:
        print("Agent not found")
    else:
        raise
except CosmosAuthError:
    print("Token expired or invalid")

async def handle_request(request):
    client = CosmosClient.from_token(
        request.headers["Authorization"].removeprefix("Bearer ").strip(),
        org_id=request.headers["X-Org-Id"],
        services=build_service_endpoints_from_env(),
    )
    try:
        return await client.search_agents({"page_size": 10})
    finally:
        await client.close()