Ai Ml Integration

"""
Embedding Utilities Template
Usage: Generate and manage embeddings for RAG, search, and similarity
"""

import os
from typing import List, Optional, Union
from dataclasses import dataclass
import numpy as np

# ===========================================
# Configuration
# ===========================================

@dataclass
class EmbeddingConfig:
    """Embedding model configuration"""
    provider: str = "openai"  # openai, google, local
    model: str = "text-embedding-3-small"
    dimensions: int = 1536
    batch_size: int = 100
    normalize: bool = True


# Provider-specific models
EMBEDDING_MODELS = {
    # OpenAI
    "text-embedding-3-small": {"provider": "openai", "dimensions": 1536},
    "text-embedding-3-large": {"provider": "openai", "dimensions": 3072},
    "text-embedding-ada-002": {"provider": "openai", "dimensions": 1536},

    # Google
    "text-embedding-004": {"provider": "google", "dimensions": 768},
    "text-multilingual-embedding-002": {"provider": "google", "dimensions": 768},

    # Local (sentence-transformers)
    "all-MiniLM-L6-v2": {"provider": "local", "dimensions": 384},
    "all-mpnet-base-v2": {"provider": "local", "dimensions": 768},
}


def get_config_from_env() -> EmbeddingConfig:
    """Load configuration from environment variables"""
    model = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
    model_info = EMBEDDING_MODELS.get(model, {})

    return EmbeddingConfig(
        provider=os.getenv("EMBEDDING_PROVIDER", model_info.get("provider", "openai")),
        model=model,
        dimensions=int(os.getenv("EMBEDDING_DIMENSIONS", model_info.get("dimensions", 1536))),
        batch_size=int(os.getenv("EMBEDDING_BATCH_SIZE", "100")),
        normalize=os.getenv("EMBEDDING_NORMALIZE", "true").lower() == "true",
    )


# ===========================================
# Embedding Client (Abstract)
# ===========================================

class EmbeddingClient:
    """Base class for embedding clients"""

    def __init__(self, config: Optional[EmbeddingConfig] = None):
        self.config = config or get_config_from_env()

    def embed(self, texts: Union[str, List[str]]) -> np.ndarray:
        """Generate embeddings for texts"""
        if isinstance(texts, str):
            texts = [texts]

        # Process in batches
        all_embeddings = []
        for i in range(0, len(texts), self.config.batch_size):
            batch = texts[i:i + self.config.batch_size]
            embeddings = self._embed_batch(batch)
            all_embeddings.extend(embeddings)

        result = np.array(all_embeddings)

        if self.config.normalize:
            result = self._normalize(result)

        return result

    def _embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Override in subclass"""
        raise NotImplementedError("Implement _embed_batch() for your provider")

    def _normalize(self, embeddings: np.ndarray) -> np.ndarray:
        """L2 normalize embeddings"""
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        return embeddings / np.maximum(norms, 1e-10)


# ===========================================
# Similarity Functions
# ===========================================

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Compute cosine similarity between two vectors"""
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))


def cosine_similarity_matrix(embeddings: np.ndarray, query: np.ndarray) -> np.ndarray:
    """Compute cosine similarity between query and all embeddings"""
    # Normalize
    embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
    query_norm = query / np.linalg.norm(query)

    return np.dot(embeddings_norm, query_norm)


def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    """Compute Euclidean distance between two vectors"""
    return float(np.linalg.norm(a - b))


def dot_product(a: np.ndarray, b: np.ndarray) -> float:
    """Compute dot product between two vectors"""
    return float(np.dot(a, b))


# ===========================================
# Search Functions
# ===========================================

@dataclass
class SearchResult:
    """Search result with score and metadata"""
    index: int
    score: float
    text: Optional[str] = None
    metadata: Optional[dict] = None


def search_similar(
    query_embedding: np.ndarray,
    embeddings: np.ndarray,
    texts: Optional[List[str]] = None,
    top_k: int = 5,
    threshold: float = 0.0,
) -> List[SearchResult]:
    """Find most similar embeddings to query"""

    # Compute similarities
    similarities = cosine_similarity_matrix(embeddings, query_embedding)

    # Get top-k indices
    top_indices = np.argsort(similarities)[::-1][:top_k]

    results = []
    for idx in top_indices:
        score = float(similarities[idx])
        if score < threshold:
            continue

        results.append(SearchResult(
            index=int(idx),
            score=score,
            text=texts[idx] if texts else None,
        ))

    return results


def deduplicate_by_similarity(
    embeddings: np.ndarray,
    threshold: float = 0.95,
) -> List[int]:
    """Remove near-duplicate embeddings, return unique indices"""

    n = len(embeddings)
    unique_indices = []

    for i in range(n):
        is_duplicate = False
        for j in unique_indices:
            sim = cosine_similarity(embeddings[i], embeddings[j])
            if sim >= threshold:
                is_duplicate = True
                break

        if not is_duplicate:
            unique_indices.append(i)

    return unique_indices


# ===========================================
# Text Chunking
# ===========================================

def chunk_text(
    text: str,
    chunk_size: int = 500,
    overlap: int = 50,
    separator: str = "\n",
) -> List[str]:
    """Split text into overlapping chunks"""

    # Split by separator first
    paragraphs = text.split(separator)

    chunks = []
    current_chunk = []
    current_size = 0

    for para in paragraphs:
        para_size = len(para)

        if current_size + para_size > chunk_size and current_chunk:
            # Save current chunk
            chunks.append(separator.join(current_chunk))

            # Keep overlap
            overlap_text = separator.join(current_chunk)
            if len(overlap_text) > overlap:
                # Find a good break point for overlap
                current_chunk = [current_chunk[-1]] if current_chunk else []
                current_size = len(current_chunk[-1]) if current_chunk else 0
            else:
                current_chunk = []
                current_size = 0

        current_chunk.append(para)
        current_size += para_size

    # Don't forget the last chunk
    if current_chunk:
        chunks.append(separator.join(current_chunk))

    return chunks


def chunk_by_tokens(
    text: str,
    max_tokens: int = 500,
    overlap_tokens: int = 50,
) -> List[str]:
    """Split text by approximate token count (4 chars ≈ 1 token)"""

    chars_per_token = 4
    chunk_size = max_tokens * chars_per_token
    overlap = overlap_tokens * chars_per_token

    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size

        # Try to break at sentence boundary
        if end < len(text):
            for boundary in [". ", "! ", "? ", "\n"]:
                last_boundary = text.rfind(boundary, start, end)
                if last_boundary > start:
                    end = last_boundary + len(boundary)
                    break

        chunks.append(text[start:end].strip())
        start = end - overlap

    return [c for c in chunks if c]  # Remove empty chunks


# ===========================================
# Caching
# ===========================================

class EmbeddingCache:
    """Simple in-memory cache for embeddings"""

    def __init__(self, max_size: int = 10000):
        self.cache: dict = {}
        self.max_size = max_size

    def get(self, text: str) -> Optional[np.ndarray]:
        """Get cached embedding"""
        return self.cache.get(hash(text))

    def set(self, text: str, embedding: np.ndarray) -> None:
        """Cache embedding"""
        if len(self.cache) >= self.max_size:
            # Simple eviction: remove first item
            first_key = next(iter(self.cache))
            del self.cache[first_key]

        self.cache[hash(text)] = embedding

    def get_or_compute(
        self,
        text: str,
        compute_fn: callable,
    ) -> np.ndarray:
        """Get from cache or compute"""
        cached = self.get(text)
        if cached is not None:
            return cached

        embedding = compute_fn(text)
        self.set(text, embedding)
        return embedding


# ===========================================
# Usage Example
# ===========================================

"""
from embedding_utils import (
    EmbeddingClient,
    get_config_from_env,
    search_similar,
    chunk_text,
)

# Initialize client
client = EmbeddingClient(get_config_from_env())

# Generate embeddings
texts = ["Hello world", "How are you?", "Machine learning is great"]
embeddings = client.embed(texts)

# Search
query_embedding = client.embed("Hi there")[0]
results = search_similar(query_embedding, embeddings, texts, top_k=2)

for r in results:
    print(f"{r.score:.3f}: {r.text}")

# Chunk long text
long_text = open("document.txt").read()
chunks = chunk_text(long_text, chunk_size=500, overlap=50)
chunk_embeddings = client.embed(chunks)
"""

/**
 * LLM Configuration Template
 * Usage: Configure multiple LLM providers with fallback
 */

// ===========================================
// Types
// ===========================================

export interface LLMConfig {
  provider: 'openai' | 'anthropic' | 'google' | 'azure' | 'local';
  model: string;
  apiKey?: string;
  baseUrl?: string;
  maxTokens?: number;
  temperature?: number;
  timeout?: number;
}

export interface LLMMessage {
  role: 'system' | 'user' | 'assistant';
  content: string;
}

export interface LLMResponse {
  content: string;
  model: string;
  usage?: {
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
  };
}

// ===========================================
// Provider Configurations
// ===========================================

export const PROVIDER_CONFIGS: Record<string, Partial<LLMConfig>> = {
  // OpenAI
  'gpt-4o': {
    provider: 'openai',
    model: 'gpt-4o',
    maxTokens: 4096,
  },
  'gpt-4o-mini': {
    provider: 'openai',
    model: 'gpt-4o-mini',
    maxTokens: 4096,
  },

  // Anthropic
  'claude-sonnet': {
    provider: 'anthropic',
    model: 'claude-sonnet-4-20250514',
    maxTokens: 8192,
  },
  'claude-haiku': {
    provider: 'anthropic',
    model: 'claude-3-5-haiku-20241022',
    maxTokens: 8192,
  },

  // Google
  'gemini-pro': {
    provider: 'google',
    model: 'gemini-2.0-flash',
    maxTokens: 8192,
  },

  // Local (Ollama)
  'llama3': {
    provider: 'local',
    model: 'llama3',
    baseUrl: 'http://localhost:11434/v1',
    maxTokens: 4096,
  },
};

// ===========================================
// Default Configuration
// ===========================================

export const DEFAULT_CONFIG: LLMConfig = {
  provider: 'openai',
  model: 'gpt-4o-mini',
  maxTokens: 4096,
  temperature: 0.7,
  timeout: 30000,
};

// ===========================================
// Environment-based Configuration
// ===========================================

export function getConfigFromEnv(): LLMConfig {
  const provider = process.env.LLM_PROVIDER as LLMConfig['provider'] || 'openai';

  const apiKeys: Record<string, string | undefined> = {
    openai: process.env.OPENAI_API_KEY,
    anthropic: process.env.ANTHROPIC_API_KEY,
    google: process.env.GOOGLE_API_KEY,
    azure: process.env.AZURE_OPENAI_API_KEY,
    local: undefined,
  };

  return {
    provider,
    model: process.env.LLM_MODEL || DEFAULT_CONFIG.model,
    apiKey: apiKeys[provider],
    baseUrl: process.env.LLM_BASE_URL,
    maxTokens: parseInt(process.env.LLM_MAX_TOKENS || '') || DEFAULT_CONFIG.maxTokens,
    temperature: parseFloat(process.env.LLM_TEMPERATURE || '') || DEFAULT_CONFIG.temperature,
    timeout: parseInt(process.env.LLM_TIMEOUT || '') || DEFAULT_CONFIG.timeout,
  };
}

// ===========================================
// LLM Client Factory
// ===========================================

export function createLLMClient(config: Partial<LLMConfig> = {}) {
  const finalConfig = { ...DEFAULT_CONFIG, ...config };

  return {
    config: finalConfig,

    async chat(messages: LLMMessage[]): Promise<LLMResponse> {
      // Implementation depends on provider
      // This is a template - implement based on your SDK choice
      throw new Error(`Implement chat() for provider: ${finalConfig.provider}`);
    },

    async complete(prompt: string): Promise<string> {
      const response = await this.chat([{ role: 'user', content: prompt }]);
      return response.content;
    },

    async stream(messages: LLMMessage[]): AsyncGenerator<string> {
      // Streaming implementation
      throw new Error(`Implement stream() for provider: ${finalConfig.provider}`);
    },
  };
}

// ===========================================
// Retry & Fallback Utilities
// ===========================================

export interface RetryConfig {
  maxRetries: number;
  initialDelay: number;
  maxDelay: number;
  backoffMultiplier: number;
}

export const DEFAULT_RETRY_CONFIG: RetryConfig = {
  maxRetries: 3,
  initialDelay: 1000,
  maxDelay: 10000,
  backoffMultiplier: 2,
};

export async function withRetry<T>(
  fn: () => Promise<T>,
  config: Partial<RetryConfig> = {}
): Promise<T> {
  const { maxRetries, initialDelay, maxDelay, backoffMultiplier } = {
    ...DEFAULT_RETRY_CONFIG,
    ...config,
  };

  let lastError: Error | undefined;
  let delay = initialDelay;

  for (let attempt = 0; attempt <= maxRetries; attempt++) {
    try {
      return await fn();
    } catch (error) {
      lastError = error as Error;

      if (attempt === maxRetries) break;

      // Check if error is retryable
      const isRetryable = isRetryableError(error);
      if (!isRetryable) throw error;

      console.warn(`Attempt ${attempt + 1} failed, retrying in ${delay}ms...`);
      await sleep(delay);
      delay = Math.min(delay * backoffMultiplier, maxDelay);
    }
  }

  throw lastError;
}

function isRetryableError(error: unknown): boolean {
  if (error instanceof Error) {
    // Rate limit errors
    if (error.message.includes('rate limit')) return true;
    if (error.message.includes('429')) return true;

    // Temporary server errors
    if (error.message.includes('500')) return true;
    if (error.message.includes('502')) return true;
    if (error.message.includes('503')) return true;

    // Network errors
    if (error.message.includes('ECONNRESET')) return true;
    if (error.message.includes('ETIMEDOUT')) return true;
  }
  return false;
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

// ===========================================
// Fallback Chain
// ===========================================

export async function withFallback<T>(
  primary: () => Promise<T>,
  fallbacks: Array<() => Promise<T>>
): Promise<T> {
  try {
    return await primary();
  } catch (primaryError) {
    console.warn('Primary LLM failed, trying fallbacks...', primaryError);

    for (let i = 0; i < fallbacks.length; i++) {
      try {
        console.log(`Trying fallback ${i + 1}...`);
        return await fallbacks[i]();
      } catch (fallbackError) {
        console.warn(`Fallback ${i + 1} failed`, fallbackError);
      }
    }

    throw new Error('All LLM providers failed');
  }
}

// ===========================================
// Usage Example
// ===========================================

/*
import { createLLMClient, getConfigFromEnv, withRetry, withFallback } from './llm-config';

// Simple usage
const client = createLLMClient(getConfigFromEnv());
const response = await client.complete('Hello, world!');

// With retry
const response = await withRetry(
  () => client.complete('Hello!'),
  { maxRetries: 3 }
);

// With fallback
const response = await withFallback(
  () => createLLMClient({ model: 'gpt-4o' }).complete('Hello!'),
  [
    () => createLLMClient({ model: 'claude-sonnet' }).complete('Hello!'),
    () => createLLMClient({ model: 'gemini-pro' }).complete('Hello!'),
  ]
);
*/