gistfile1.txt

import os
from typing import Any, Dict, List

import requests

# Get access token from environment variables
access_token = os.environ.get("DASH_TOKEN")
if not access_token:
    raise ValueError("Please set the DASH_TOKEN environment variable")

# API base URL - adjust if needed
API_BASE_URL = "https://dash-api.orionarm.ai/api"

# Common headers for all API requests
headers = {
    "accept": "application/json",
    "Authorization": f"Bearer {access_token}",
    "Content-Type": "application/json",
}


def call_embedding_api(
    texts: List[str], model: str = None, dimensions: int = None
) -> Dict[str, Any]:
    """
    Call the embedding API to generate embeddings for the given texts.

    Args:
        texts: List of text strings to generate embeddings for
        model: Optional embedding model name (defaults to server's default model)
        dimensions: Optional embedding dimensions

    Returns:
        API response containing embeddings
    """
    payload = {"texts": texts, "timeout": 60, "enable_cache": True}

    # Add optional parameters if provided
    if model:
        payload["model"] = model
    if dimensions:
        payload["dimensions"] = dimensions

    response = requests.post(
        f"{API_BASE_URL}/embeddings",
        headers=headers,
        json=payload,
    )

    if response.status_code == 200:
        return response.json()
    else:
        print(f"Error calling embedding API: {response.status_code}")
        print(response.text)
        return None


def call_rerank_api(
    query: str, documents: List[Dict], top_k: int = 3
) -> Dict[str, Any]:
    """
    Call the rerank API to reorder documents based on relevance to query.

    Args:
        query: The query text to rank documents against
        documents: List of document dictionaries, each with id, content, and optional relevance_score
        top_k: Number of top results to return

    Returns:
        API response containing reranked documents
    """
    # Format documents for the API
    formatted_documents = [
        {
            "id": doc.get("id", i),
            "content": doc["content"],
            "relevance_score": doc.get("relevance_score", 0.0),
        }
        for i, doc in enumerate(documents)
    ]

    payload = {
        "query": query,
        "documents": formatted_documents,
        "top_k": top_k,
        "timeout": 30,
        "chunk_size": 200,
    }

    response = requests.post(
        f"{API_BASE_URL}/rerank",
        headers=headers,
        json=payload,
    )

    if response.status_code == 200:
        return response.json()
    else:
        print(f"Error calling rerank API: {response.status_code}")
        print(response.text)
        return None


def demo_embedding_api():
    """Demo the embedding API with sample texts"""
    print("\n=== EMBEDDING API DEMO ===")

    # Sample texts to embed
    texts = [
        "The quick brown fox jumps over the lazy dog",
        "Machine learning models can generate text embeddings",
        "Embeddings are useful for semantic search and clustering",
    ]

    print(f"Generating embeddings for {len(texts)} texts...")
    result = call_embedding_api(texts)

    if result:
        print(f"Successfully generated embeddings using model: {result['model']}")
        print(f"Embedding dimensions: {len(result['embeddings'][0])}")

        # Print a sample of the first embedding vector
        first_embedding = result["embeddings"][0]
        print(f"Sample of first embedding vector: {first_embedding[:5]}...")


def demo_rerank_api():
    """Demo the rerank API with sample query and documents"""
    print("\n=== RERANK API DEMO ===")

    # Sample query
    query = "artificial intelligence applications"

    # Sample documents
    documents = [
        {
            "id": 1,
            "content": "AI is revolutionizing healthcare through improved diagnostics and personalized treatment plans.",
        },
        {
            "id": 2,
            "content": "The golden retriever is a medium-large gun dog that was bred to retrieve shot waterfowl.",
        },
        {
            "id": 3,
            "content": "Machine learning applications in finance help detect fraud and automate trading strategies.",
        },
        {
            "id": 4,
            "content": "Neural networks are a fundamental component of modern artificial intelligence systems.",
        },
        {
            "id": 5,
            "content": "The history of artificial intelligence dates back to ancient myths and stories.",
        },
    ]

    print(f"Reranking {len(documents)} documents based on query: '{query}'")
    result = call_rerank_api(query, documents, top_k=3)

    if result:
        print("\nReranked results:")
        for i, doc in enumerate(result["results"]):
            print(
                f"{i+1}. [ID: {doc['id']}, Score: {doc['relevance_score']:.4f}] {doc['content']}"
            )


if __name__ == "__main__":
    # Run both demos
    demo_embedding_api()
    demo_rerank_api()
    print("\nDemo completed successfully!")