aetherroute 0.1.0__py3-none-any.whl

aetherroute/__init__.py

@@ -0,0 +1,8 @@
+"""
+AetherRoute - Multi-provider LLM orchestration and routing framework.
+"""
+
+from aetherroute.config import load_config, AetherRouteConfig
+
+__version__ = "0.1.0"
+__all__ = ["load_config", "AetherRouteConfig"]

aetherroute/adapters/__init__.py

@@ -0,0 +1,8 @@
+from aetherroute.adapters.prompt import PromptAdapter
+from aetherroute.adapters.token_counter import count_tokens, count_messages_tokens
+
+__all__ = [
+    "PromptAdapter",
+    "count_tokens",
+    "count_messages_tokens"
+]

aetherroute/adapters/prompt.py

@@ -0,0 +1,129 @@
+import re
+import logging
+from typing import List, Dict, Any, Optional
+from aetherroute.adapters.token_counter import count_messages_tokens, count_tokens
+
+logger = logging.getLogger("aetherroute.adapters.prompt")
+
+class PromptAdapter:
+    """
+    Normalizes templates and message histories for specific LLM providers.
+    Trims history to fit within a provider's max context limit.
+    """
+    
+    @staticmethod
+    def render_template(template: str, variables: Dict[str, Any]) -> str:
+        """
+        Renders variables into double curly-braces template placeholders (e.g. {{variable}}).
+        """
+        rendered = template
+        for key, val in variables.items():
+            pattern = re.compile(r"\{\{\s*" + re.escape(key) + r"\s*\}\}")
+            rendered = pattern.sub(str(val), rendered)
+        return rendered
+
+    @classmethod
+    def normalize_messages(
+        cls, 
+        messages: List[Dict[str, str]], 
+        provider_name: str,
+        max_context: int,
+        reserved_output_tokens: int = 2048
+    ) -> List[Dict[str, str]]:
+        """
+        Normalizes role types, alternates roles if necessary (Anthropic),
+        and trims historical messages to fit the provider's max context budget.
+        """
+        if not messages:
+            return []
+
+        # 1. Filter out system messages and other messages
+        system_msgs = [m for m in messages if m.get("role") == "system"]
+        other_msgs = [m for m in messages if m.get("role") != "system"]
+
+        # 2. Provider specific adjustments
+        normalized_others = []
+        if provider_name == "anthropic":
+            # Anthropic Claude requires alternating user and assistant messages, starting with user.
+            # Merge consecutive identical roles, and drop/convert initial assistant messages.
+            current_role = None
+            current_content = []
+
+            for msg in other_msgs:
+                role = msg.get("role", "user")
+                # Normalize roles
+                if role not in ["user", "assistant"]:
+                    role = "user"
+
+                if current_role is None:
+                    if role == "assistant":
+                        # Cannot start with assistant, convert to user or ignore
+                        # Let's convert it to user so we don't lose context
+                        role = "user"
+                    current_role = role
+                    current_content.append(msg.get("content", ""))
+                elif role == current_role:
+                    # Merge consecutive identical roles
+                    current_content.append(msg.get("content", ""))
+                else:
+                    # Save previous
+                    normalized_others.append({
+                        "role": current_role,
+                        "content": "\n\n".join(current_content)
+                    })
+                    current_role = role
+                    current_content = [msg.get("content", "")]
+            
+            # Add the last message group
+            if current_role:
+                normalized_others.append({
+                    "role": current_role,
+                    "content": "\n\n".join(current_content)
+                })
+        else:
+            # OpenAI, Mistral, Ollama accept standard system/user/assistant messages.
+            for msg in other_msgs:
+                role = msg.get("role", "user")
+                if role not in ["user", "assistant", "system"]:
+                    role = "user"
+                normalized_others.append({
+                    "role": role,
+                    "content": msg.get("content", "")
+                })
+
+        # 3. Handle context window pruning (sliding window over chat history)
+        # We always want to keep the system messages if they exist.
+        # We prune from the oldest of other_msgs (index 0 of normalized_others)
+        allowed_tokens = max_context - reserved_output_tokens
+        
+        system_tokens = count_messages_tokens(system_msgs)
+        if system_tokens > allowed_tokens:
+            # System message itself exceeds the allowed size! Prune system message characters as fallback.
+            logger.warning("System message exceeds context budget. Truncating system message.")
+            for s_msg in system_msgs:
+                s_msg["content"] = s_msg["content"][:allowed_tokens * 4]
+            return system_msgs
+
+        # Prune older user/assistant messages until we fit
+        final_messages = system_msgs + normalized_others
+        total_tokens = count_messages_tokens(final_messages)
+        
+        while total_tokens > allowed_tokens and len(normalized_others) > 1:
+            # Remove the oldest history message
+            removed = normalized_others.pop(0)
+            logger.info(f"Context limit reached. Pruning message with role: {removed['role']}")
+            final_messages = system_msgs + normalized_others
+            total_tokens = count_messages_tokens(final_messages)
+
+        # If it still exceeds (e.g. only one huge query message remaining), truncate the text content of that message
+        if total_tokens > allowed_tokens and normalized_others:
+            msg_to_truncate = normalized_others[0]
+            current_tokens = count_tokens(msg_to_truncate["content"])
+            excess_tokens = total_tokens - allowed_tokens
+            allowed_msg_tokens = max(100, current_tokens - excess_tokens)
+            
+            # Simple truncate by character approximation
+            msg_to_truncate["content"] = msg_to_truncate["content"][:allowed_msg_tokens * 4] + "... [truncated]"
+            final_messages = system_msgs + normalized_others
+
+        return final_messages

aetherroute/adapters/token_counter.py

@@ -0,0 +1,44 @@
+import logging
+from typing import List, Dict, Any
+
+logger = logging.getLogger("aetherroute.adapters.token_counter")
+
+try:
+    import tiktoken
+    _TIKTOKEN_AVAILABLE = True
+except ImportError:
+    _TIKTOKEN_AVAILABLE = False
+    logger.warning("tiktoken library not found. Falling back to rough character estimation (1 token ≈ 4 characters).")
+
+def count_tokens(text: str, model_name: str = "gpt-4o-mini") -> int:
+    """
+    Returns the token count of a given string.
+    Uses tiktoken if available, otherwise falls back to character division.
+    """
+    if not text:
+        return 0
+        
+    if _TIKTOKEN_AVAILABLE:
+        try:
+            # Try to get the encoding for the model
+            try:
+                encoding = tiktoken.encoding_for_model(model_name)
+            except KeyError:
+                encoding = tiktoken.get_encoding("cl100k_base")
+            return len(encoding.encode(text))
+        except Exception as e:
+            logger.debug(f"Error counting tokens with tiktoken: {e}. Falling back to character estimation.")
+            
+    # Fallback: 1 token is roughly 4 characters on average for English text
+    return max(1, len(text) // 4)
+
+def count_messages_tokens(messages: List[Dict[str, str]], model_name: str = "gpt-4o-mini") -> int:
+    """
+    Estimates the total token count of a message list.
+    """
+    token_sum = 0
+    for msg in messages:
+        token_sum += count_tokens(msg.get("content", ""), model_name)
+        # Add a few overhead tokens per message (role, structure)
+        token_sum += 4
+    return token_sum + 3 # overall conversation overhead

aetherroute/cache/__init__.py

@@ -0,0 +1,3 @@
+from aetherroute.cache.semantic import SemanticCache
+
+__all__ = ["SemanticCache"]

aetherroute/cache/semantic.py

@@ -0,0 +1,153 @@
+import json
+import logging
+import asyncio
+from typing import Dict, Any, Optional, List, Tuple
+from aetherroute.context.curator import ContextCurator
+
+logger = logging.getLogger("aetherroute.cache.semantic")
+
+try:
+    import redis.asyncio as aioredis
+    _REDIS_AVAILABLE = True
+except ImportError:
+    _REDIS_AVAILABLE = False
+    logger.warning("redis-py not installed. Caching will run strictly in-memory.")
+
+class SemanticCache:
+    """
+    Semantic and exact-match cache for LLM completions.
+    Tries to connect to Redis, and automatically falls back to an in-memory database.
+    Matches queries fuzzy-style using TF-IDF cosine similarity.
+    """
+    def __init__(self, config: Any):
+        self.config = config
+        self.redis_url = getattr(config.cache, "redis_url", "redis://localhost:6379")
+        self.threshold = getattr(config.cache, "semantic_threshold", 0.85)
+        self.ttl = getattr(config.cache, "ttl_seconds", 3600)
+        self.client: Optional[Any] = None
+        self.use_redis = False
+        
+        # Local in-memory cache fallback: maps query -> (response_dict, expiry_timestamp)
+        self._in_memory_cache: Dict[str, Tuple[Dict[str, Any], float]] = {}
+
+    async def connect(self) -> None:
+        """Attempts to connect to Redis. Falls back to in-memory on failure."""
+        if _REDIS_AVAILABLE:
+            try:
+                self.client = aioredis.from_url(self.redis_url, decode_responses=True)
+                # Test connection
+                await asyncio.wait_for(self.client.ping(), timeout=2.0)
+                self.use_redis = True
+                logger.info(f"Connected to Redis semantic cache at {self.redis_url}")
+                return
+            except Exception as e:
+                logger.warning(f"Failed to connect to Redis ({e}). Falling back to In-Memory Cache.")
+        
+        self.use_redis = False
+        self.client = None
+
+    async def get(self, query: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieves a cached response if an exact or semantically similar query is found.
+        """
+        # Ensure connection state is checked
+        if self.use_redis and self.client is None:
+            await self.connect()
+
+        cleaned_query = query.strip()
+        
+        if self.use_redis and self.client:
+            try:
+                # 1. Check exact match
+                exact_key = f"cache:exact:{cleaned_query}"
+                exact_val = await self.client.get(exact_key)
+                if exact_val:
+                    logger.info("SemanticCache: Exact match found in Redis.")
+                    return json.loads(exact_val)
+
+                # 2. Check semantic match
+                # Get all cached queries
+                keys = await self.client.keys("cache:exact:*")
+                for key in keys:
+                    original_query = key.replace("cache:exact:", "", 1)
+                    similarity = ContextCurator.calculate_tfidf_similarity(cleaned_query, original_query)
+                    if similarity >= self.threshold:
+                        logger.info(f"SemanticCache: Fuzzy match found in Redis (similarity={similarity:.4f}).")
+                        val = await self.client.get(key)
+                        if val:
+                            return json.loads(val)
+            except Exception as e:
+                logger.error(f"Redis cache read error: {e}. Falling back to local search.")
+                self.use_redis = False
+
+        # In-memory search (either fallback or primary)
+        import time
+        now = time.time()
+        
+        # 1. Exact match in memory
+        if cleaned_query in self._in_memory_cache:
+            val, expiry = self._in_memory_cache[cleaned_query]
+            if now < expiry:
+                logger.info("SemanticCache: Exact match found in memory.")
+                return val
+            else:
+                del self._in_memory_cache[cleaned_query]
+
+        # 2. Semantic match in memory
+        for orig_query, (val, expiry) in list(self._in_memory_cache.items()):
+            if now < expiry:
+                similarity = ContextCurator.calculate_tfidf_similarity(cleaned_query, orig_query)
+                if similarity >= self.threshold:
+                    logger.info(f"SemanticCache: Fuzzy match found in memory (similarity={similarity:.4f}).")
+                    return val
+            else:
+                del self._in_memory_cache[orig_query]
+
+        return None
+
+    async def set(self, query: str, response: Dict[str, Any], ttl: Optional[int] = None) -> None:
+        """
+        Caches a query response.
+        """
+        if self.use_redis and self.client is None:
+            await self.connect()
+
+        cleaned_query = query.strip()
+        cache_ttl = ttl if ttl is not None else self.ttl
+        response_str = json.dumps(response)
+
+        if self.use_redis and self.client:
+            try:
+                exact_key = f"cache:exact:{cleaned_query}"
+                await self.client.set(exact_key, response_str, ex=cache_ttl)
+                logger.debug(f"Cached query in Redis with TTL={cache_ttl}")
+                return
+            except Exception as e:
+                logger.error(f"Redis cache write error: {e}. Writing to in-memory fallback.")
+                self.use_redis = False
+
+        # Write to memory fallback
+        import time
+        expiry = time.time() + cache_ttl
+        self._in_memory_cache[cleaned_query] = (response, expiry)
+        logger.debug(f"Cached query in-memory with TTL={cache_ttl}")
+
+    async def clear(self) -> None:
+        """Clears all cached entries."""
+        self._in_memory_cache.clear()
+        if self.use_redis and self.client:
+            try:
+                keys = await self.client.keys("cache:*")
+                if keys:
+                    await self.client.delete(*keys)
+            except Exception as e:
+                logger.error(f"Failed to clear Redis cache: {e}")
+                
+    async def close(self) -> None:
+        """Closes Redis connection if open."""
+        if self.client:
+            try:
+                await self.client.aclose()
+            except Exception:
+                pass
+            self.client = None

aetherroute/config.py

@@ -0,0 +1,65 @@
+import os
+import yaml
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+from pydantic import BaseModel, Field
+
+class ModelConfig(BaseModel):
+    name: str
+    max_context: int
+    cost_per_1k_input: float
+    cost_per_1k_output: float
+    latency_sla: float
+
+class ProviderConfig(BaseModel):
+    api_key_env: Optional[str] = None
+    base_url: Optional[str] = None
+    models: List[ModelConfig]
+
+class TaskWeights(BaseModel):
+    accuracy: float
+    cost: float
+    latency: float
+
+class RoutingConfig(BaseModel):
+    default_model: str
+    cost_ceiling_request: float
+    cost_ceiling_session: float
+    historical_db_path: str
+    task_weights: Dict[str, TaskWeights]
+
+class CacheConfig(BaseModel):
+    redis_url: str
+    semantic_threshold: float
+    ttl_seconds: int
+
+class SecurityConfig(BaseModel):
+    prompt_injection_action: str
+    permissions: Dict[str, Dict[str, List[str]]]
+
+class AetherRouteConfig(BaseModel):
+    providers: Dict[str, ProviderConfig]
+    routing: RoutingConfig
+    cache: CacheConfig
+    security: SecurityConfig
+
+def load_config(config_path: Optional[str] = None) -> AetherRouteConfig:
+    """Load configuration from a YAML file. Looks in default locations if path not provided."""
+    if config_path is None:
+        # Check current working directory, then package root parent
+        paths_to_check = [
+            Path.cwd() / "config.yaml",
+            Path(__file__).parent.parent / "config.yaml",
+        ]
+        for p in paths_to_check:
+            if p.exists():
+                config_path = str(p)
+                break
+        
+    if not config_path or not Path(config_path).exists():
+        raise FileNotFoundError(f"Configuration file not found. Checked default locations. Please provide a valid config_path.")
+        
+    with open(config_path, "r") as f:
+        data = yaml.safe_load(f)
+        
+    return AetherRouteConfig(**data)

aetherroute/context/__init__.py

@@ -0,0 +1,3 @@
+from aetherroute.context.curator import ContextCurator
+
+__all__ = ["ContextCurator"]

aetherroute/context/curator.py

@@ -0,0 +1,164 @@
+import re
+import math
+import logging
+from typing import List, Dict, Any, Optional
+from collections import Counter
+from aetherroute.adapters.token_counter import count_tokens, count_messages_tokens
+
+logger = logging.getLogger("aetherroute.context.curator")
+
+class ContextCurator:
+    """
+    Ranks context documents based on TF-IDF relevance to the query, prunes context
+    to fit token budgets, and uses a cheap provider/model to summarize older conversation history.
+    """
+    def __init__(self, config: Any):
+        self.config = config
+
+    @staticmethod
+    def _tokenize(text: str) -> List[str]:
+        """Simple lowercase word tokenization."""
+        return re.findall(r'[a-z0-9]+', text.lower())
+
+    @classmethod
+    def calculate_tfidf_similarity(cls, query: str, document: str) -> float:
+        """
+        Calculates cosine similarity between query and document using simple TF-IDF vectors.
+        """
+        q_tokens = cls._tokenize(query)
+        d_tokens = cls._tokenize(document)
+        
+        if not q_tokens or not d_tokens:
+            return 0.0
+            
+        q_counts = Counter(q_tokens)
+        d_counts = Counter(d_tokens)
+        
+        # Calculate term frequencies
+        q_tf = {word: count / len(q_tokens) for word, count in q_counts.items()}
+        d_tf = {word: count / len(d_tokens) for word, count in d_counts.items()}
+        
+        # Calculate cosine similarity
+        dot_product = 0.0
+        for word in q_tf:
+            if word in d_tf:
+                dot_product += q_tf[word] * d_tf[word]
+                
+        q_norm = math.sqrt(sum(val ** 2 for val in q_tf.values()))
+        d_norm = math.sqrt(sum(val ** 2 for val in d_tf.values()))
+        
+        if q_norm == 0 or d_norm == 0:
+            return 0.0
+            
+        return dot_product / (q_norm * d_norm)
+
+    def rank_and_prune_documents(
+        self, 
+        query: str, 
+        documents: List[Dict[str, Any]], 
+        max_tokens: int,
+        model_name: str = "gpt-4o-mini"
+    ) -> List[Dict[str, Any]]:
+        """
+        Ranks documents based on query similarity and includes as many high-scoring 
+        documents as possible within the specified max_tokens limit.
+        
+        Args:
+            query: The user's search query.
+            documents: List of dicts, e.g. [{"id": 1, "text": "content..."}, ...]
+            max_tokens: Limit on how many tokens of context we can include.
+            model_name: Model used to compute token count.
+            
+        Returns:
+            Pruned and sorted list of relevant documents.
+        """
+        scored_docs = []
+        for doc in documents:
+            text = doc.get("text", "")
+            score = self.calculate_tfidf_similarity(query, text)
+            # Store doc with its score
+            scored_docs.append((doc, score))
+            
+        # Sort by score descending
+        scored_docs.sort(key=lambda x: x[1], reverse=True)
+        
+        selected_docs = []
+        accumulated_tokens = 0
+        
+        for doc, score in scored_docs:
+            doc_tokens = count_tokens(doc.get("text", ""), model_name)
+            if accumulated_tokens + doc_tokens <= max_tokens:
+                selected_docs.append(doc)
+                accumulated_tokens += doc_tokens
+            else:
+                logger.info(f"Pruning document (score={score:.4f}) to stay under token budget of {max_tokens}.")
+                
+        return selected_docs
+
+    async def summarize_history(
+        self, 
+        messages: List[Dict[str, str]], 
+        cheap_provider: Any, 
+        cheap_model: str,
+        keep_last_n: int = 3
+    ) -> List[Dict[str, str]]:
+        """
+        Compresses older messages in a conversation history by summarizing them
+        using a cheap provider model, keeping the last N messages as active context.
+        """
+        # If there aren't enough messages to warrant summarization, return as is
+        system_msgs = [m for m in messages if m.get("role") == "system"]
+        chat_msgs = [m for m in messages if m.get("role") != "system"]
+        
+        if len(chat_msgs) <= keep_last_n + 1:
+            return messages
+
+        # Split into old history and fresh messages
+        old_history = chat_msgs[:-keep_last_n]
+        fresh_messages = chat_msgs[-keep_last_n:]
+
+        # Build a text compilation of the old history
+        formatted_history = []
+        for msg in old_history:
+            role = msg.get("role", "user").upper()
+            content = msg.get("content", "")
+            formatted_history.append(f"{role}: {content}")
+            
+        history_text = "\n".join(formatted_history)
+        
+        # Prepare the summarization query
+        summarize_prompt = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant. Summarize the following chat conversation highlights in a concise, dense paragraph. Focus on user requests, key decisions made, and assistant answers."
+            },
+            {
+                "role": "user",
+                "content": f"Conversation to summarize:\n\n{history_text}"
+            }
+        ]
+        
+        logger.info(f"Summarizing {len(old_history)} messages using cheap model '{cheap_model}'...")
+        
+        try:
+            # Generate the summary using the cheap provider
+            summary_response = await cheap_provider.generate(
+                messages=summarize_prompt,
+                model=cheap_model,
+                options={"max_tokens": 500, "temperature": 0.3}
+            )
+            summary_text = summary_response.get("text", "Error generating summary.")
+            
+            # Form the new message history
+            new_history = system_msgs.copy()
+            new_history.append({
+                "role": "system",
+                "content": f"Summary of earlier conversation highlights: {summary_text}"
+            })
+            new_history.extend(fresh_messages)
+            
+            logger.info("Successfully summarized older chat history.")
+            return new_history
+        except Exception as e:
+            logger.error(f"Failed to summarize older context: {e}. Returning original messages.")
+            return messages

aetherroute/cost/__init__.py

@@ -0,0 +1,3 @@
+from aetherroute.cost.governor import CostGovernor, CostLimitExceeded
+
+__all__ = ["CostGovernor", "CostLimitExceeded"]